End of training
Browse files- checkpoint-1000/latest +1 -0
- checkpoint-1000/pytorch_model/mp_rank_00_model_states.pt +3 -0
- checkpoint-1000/pytorch_model/zero_pp_rank_0_mp_rank_00_optim_states.pt +3 -0
- checkpoint-1000/random_states_0.pkl +3 -0
- checkpoint-1000/scheduler.bin +3 -0
- checkpoint-1000/unet/config.json +58 -0
- checkpoint-1000/unet/diffusion_pytorch_model.bin +3 -0
- checkpoint-1000/zero_to_fp32.py +461 -0
- checkpoint-500/latest +1 -0
- checkpoint-500/pytorch_model/mp_rank_00_model_states.pt +3 -0
- checkpoint-500/pytorch_model/zero_pp_rank_0_mp_rank_00_optim_states.pt +3 -0
- checkpoint-500/random_states_0.pkl +3 -0
- checkpoint-500/scheduler.bin +3 -0
- checkpoint-500/unet/config.json +58 -0
- checkpoint-500/unet/diffusion_pytorch_model.bin +3 -0
- checkpoint-500/zero_to_fp32.py +461 -0
- feature_extractor/preprocessor_config.json +28 -0
- model_index.json +33 -0
- safety_checker/config.json +168 -0
- safety_checker/pytorch_model.bin +3 -0
- scheduler/scheduler_config.json +14 -0
- text_encoder/config.json +25 -0
- text_encoder/pytorch_model.bin +3 -0
- tokenizer/merges.txt +0 -0
- tokenizer/special_tokens_map.json +24 -0
- tokenizer/tokenizer_config.json +33 -0
- tokenizer/vocab.json +0 -0
- unet/config.json +58 -0
- unet/diffusion_pytorch_model.bin +3 -0
- vae/config.json +31 -0
- vae/diffusion_pytorch_model.bin +3 -0
checkpoint-1000/latest
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
pytorch_model
|
checkpoint-1000/pytorch_model/mp_rank_00_model_states.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:608cf70ca94237fee5d8f51ddd99fcc599e3d843a79caa73a5902320db8e151c
|
| 3 |
+
size 1719248603
|
checkpoint-1000/pytorch_model/zero_pp_rank_0_mp_rank_00_optim_states.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:fd0d40315512a5e3c89cf55e50833faa41de3014f04fd0701c66c9a0c64c8f63
|
| 3 |
+
size 10314315422
|
checkpoint-1000/random_states_0.pkl
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:89d6cf50fbfbbb5ad8ab1f000ff7a0b61e24ac197c2d581b68f5f7474b41b9fd
|
| 3 |
+
size 14599
|
checkpoint-1000/scheduler.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:9f6a2f29716edbefad84c26b160fec83650b64f7a7d87713c4632755f7c3595b
|
| 3 |
+
size 563
|
checkpoint-1000/unet/config.json
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_class_name": "UNet2DConditionModel",
|
| 3 |
+
"_diffusers_version": "0.16.0.dev0",
|
| 4 |
+
"_name_or_path": "runwayml/stable-diffusion-v1-5",
|
| 5 |
+
"act_fn": "silu",
|
| 6 |
+
"attention_head_dim": 8,
|
| 7 |
+
"block_out_channels": [
|
| 8 |
+
320,
|
| 9 |
+
640,
|
| 10 |
+
1280,
|
| 11 |
+
1280
|
| 12 |
+
],
|
| 13 |
+
"center_input_sample": false,
|
| 14 |
+
"class_embed_type": null,
|
| 15 |
+
"class_embeddings_concat": false,
|
| 16 |
+
"conv_in_kernel": 3,
|
| 17 |
+
"conv_out_kernel": 3,
|
| 18 |
+
"cross_attention_dim": 768,
|
| 19 |
+
"cross_attention_norm": null,
|
| 20 |
+
"down_block_types": [
|
| 21 |
+
"CrossAttnDownBlock2D",
|
| 22 |
+
"CrossAttnDownBlock2D",
|
| 23 |
+
"CrossAttnDownBlock2D",
|
| 24 |
+
"DownBlock2D"
|
| 25 |
+
],
|
| 26 |
+
"downsample_padding": 1,
|
| 27 |
+
"dual_cross_attention": false,
|
| 28 |
+
"encoder_hid_dim": null,
|
| 29 |
+
"flip_sin_to_cos": true,
|
| 30 |
+
"freq_shift": 0,
|
| 31 |
+
"in_channels": 4,
|
| 32 |
+
"layers_per_block": 2,
|
| 33 |
+
"mid_block_only_cross_attention": null,
|
| 34 |
+
"mid_block_scale_factor": 1,
|
| 35 |
+
"mid_block_type": "UNetMidBlock2DCrossAttn",
|
| 36 |
+
"norm_eps": 1e-05,
|
| 37 |
+
"norm_num_groups": 32,
|
| 38 |
+
"num_class_embeds": null,
|
| 39 |
+
"only_cross_attention": false,
|
| 40 |
+
"out_channels": 4,
|
| 41 |
+
"projection_class_embeddings_input_dim": null,
|
| 42 |
+
"resnet_out_scale_factor": 1.0,
|
| 43 |
+
"resnet_skip_time_act": false,
|
| 44 |
+
"resnet_time_scale_shift": "default",
|
| 45 |
+
"sample_size": 64,
|
| 46 |
+
"time_cond_proj_dim": null,
|
| 47 |
+
"time_embedding_act_fn": null,
|
| 48 |
+
"time_embedding_type": "positional",
|
| 49 |
+
"timestep_post_act": null,
|
| 50 |
+
"up_block_types": [
|
| 51 |
+
"UpBlock2D",
|
| 52 |
+
"CrossAttnUpBlock2D",
|
| 53 |
+
"CrossAttnUpBlock2D",
|
| 54 |
+
"CrossAttnUpBlock2D"
|
| 55 |
+
],
|
| 56 |
+
"upcast_attention": false,
|
| 57 |
+
"use_linear_projection": false
|
| 58 |
+
}
|
checkpoint-1000/unet/diffusion_pytorch_model.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:056230474fd49cfd202d49386a87ff77e7333357bfdbdb43b6621c4b4a1367cb
|
| 3 |
+
size 1719188507
|
checkpoint-1000/zero_to_fp32.py
ADDED
|
@@ -0,0 +1,461 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python
|
| 2 |
+
|
| 3 |
+
# Copyright (c) Microsoft Corporation.
|
| 4 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 5 |
+
|
| 6 |
+
# DeepSpeed Team
|
| 7 |
+
|
| 8 |
+
# This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
|
| 9 |
+
# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
|
| 10 |
+
# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
|
| 11 |
+
# application.
|
| 12 |
+
#
|
| 13 |
+
# example: python zero_to_fp32.py . pytorch_model.bin
|
| 14 |
+
|
| 15 |
+
import argparse
|
| 16 |
+
import torch
|
| 17 |
+
import glob
|
| 18 |
+
import math
|
| 19 |
+
import os
|
| 20 |
+
import re
|
| 21 |
+
from collections import OrderedDict
|
| 22 |
+
|
| 23 |
+
# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
|
| 24 |
+
# DeepSpeed data structures it has to be available in the current python environment.
|
| 25 |
+
from deepspeed.utils import logger
|
| 26 |
+
from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
|
| 27 |
+
FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES)
|
| 28 |
+
|
| 29 |
+
debug = 0
|
| 30 |
+
|
| 31 |
+
# load to cpu
|
| 32 |
+
device = torch.device('cpu')
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def atoi(text):
|
| 36 |
+
return int(text) if text.isdigit() else text
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def natural_keys(text):
|
| 40 |
+
'''
|
| 41 |
+
alist.sort(key=natural_keys) sorts in human order
|
| 42 |
+
http://nedbatchelder.com/blog/200712/human_sorting.html
|
| 43 |
+
(See Toothy's implementation in the comments)
|
| 44 |
+
'''
|
| 45 |
+
return [atoi(c) for c in re.split(r'(\d+)', text)]
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def get_model_state_file(checkpoint_dir, zero_stage):
|
| 49 |
+
if not os.path.isdir(checkpoint_dir):
|
| 50 |
+
raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
|
| 51 |
+
|
| 52 |
+
# there should be only one file
|
| 53 |
+
if zero_stage == 2:
|
| 54 |
+
file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
|
| 55 |
+
elif zero_stage == 3:
|
| 56 |
+
file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
|
| 57 |
+
|
| 58 |
+
if not os.path.exists(file):
|
| 59 |
+
raise FileNotFoundError(f"can't find model states file at '{file}'")
|
| 60 |
+
|
| 61 |
+
return file
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def get_optim_files(checkpoint_dir):
|
| 65 |
+
# XXX: need to test that this simple glob rule works for multi-node setup too
|
| 66 |
+
optim_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*_optim_states.pt")), key=natural_keys)
|
| 67 |
+
|
| 68 |
+
if len(optim_files) == 0:
|
| 69 |
+
raise FileNotFoundError(f"can't find '*_optim_states.pt' files in directory '{checkpoint_dir}'")
|
| 70 |
+
|
| 71 |
+
return optim_files
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def parse_model_state(file):
|
| 75 |
+
state_dict = torch.load(file, map_location=device)
|
| 76 |
+
|
| 77 |
+
if BUFFER_NAMES not in state_dict:
|
| 78 |
+
raise ValueError(f"{file} is not a model state checkpoint")
|
| 79 |
+
buffer_names = state_dict[BUFFER_NAMES]
|
| 80 |
+
if debug:
|
| 81 |
+
print("Found buffers:", buffer_names)
|
| 82 |
+
|
| 83 |
+
# recover just the buffers while restoring them to fp32 if they were saved in fp16
|
| 84 |
+
buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
|
| 85 |
+
param_shapes = state_dict[PARAM_SHAPES]
|
| 86 |
+
|
| 87 |
+
# collect parameters that are included in param_shapes
|
| 88 |
+
param_names = []
|
| 89 |
+
for s in param_shapes:
|
| 90 |
+
for name in s.keys():
|
| 91 |
+
param_names.append(name)
|
| 92 |
+
|
| 93 |
+
# record shared parameters so that they can be recovered based on partners
|
| 94 |
+
# this is because such parameters holding reference only are not saved by optimizer
|
| 95 |
+
shared_params = []
|
| 96 |
+
for param in state_dict["module"]:
|
| 97 |
+
if param not in [*param_names, *buffer_names]:
|
| 98 |
+
for share_param in state_dict["module"]:
|
| 99 |
+
if (state_dict["module"][share_param].data_ptr() == state_dict["module"][param].data_ptr()
|
| 100 |
+
and share_param != param):
|
| 101 |
+
shared_params.append([param, share_param])
|
| 102 |
+
break
|
| 103 |
+
|
| 104 |
+
ds_version = state_dict.get(DS_VERSION, None)
|
| 105 |
+
|
| 106 |
+
return buffers, param_shapes, shared_params, ds_version
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def parse_optim_states(files, ds_checkpoint_dir):
|
| 110 |
+
|
| 111 |
+
total_files = len(files)
|
| 112 |
+
state_dicts = []
|
| 113 |
+
for f in files:
|
| 114 |
+
state_dicts.append(torch.load(f, map_location=device))
|
| 115 |
+
|
| 116 |
+
if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
|
| 117 |
+
raise ValueError(f"{files[0]} is not a zero checkpoint")
|
| 118 |
+
zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
|
| 119 |
+
world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
|
| 120 |
+
|
| 121 |
+
# For ZeRO-2 each param group can have different partition_count as data parallelism for expert
|
| 122 |
+
# parameters can be different from data parallelism for non-expert parameters. So we can just
|
| 123 |
+
# use the max of the partition_count to get the dp world_size.
|
| 124 |
+
|
| 125 |
+
if type(world_size) is list:
|
| 126 |
+
world_size = max(world_size)
|
| 127 |
+
|
| 128 |
+
if world_size != total_files:
|
| 129 |
+
raise ValueError(
|
| 130 |
+
f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
|
| 131 |
+
"Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
|
| 132 |
+
)
|
| 133 |
+
|
| 134 |
+
# the groups are named differently in each stage
|
| 135 |
+
if zero_stage == 2:
|
| 136 |
+
fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
|
| 137 |
+
elif zero_stage == 3:
|
| 138 |
+
fp32_groups_key = FP32_FLAT_GROUPS
|
| 139 |
+
else:
|
| 140 |
+
raise ValueError(f"unknown zero stage {zero_stage}")
|
| 141 |
+
|
| 142 |
+
if zero_stage == 2:
|
| 143 |
+
fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
|
| 144 |
+
elif zero_stage == 3:
|
| 145 |
+
# if there is more than one param group, there will be multiple flattened tensors - one
|
| 146 |
+
# flattened tensor per group - for simplicity merge them into a single tensor
|
| 147 |
+
#
|
| 148 |
+
# XXX: could make the script more memory efficient for when there are multiple groups - it
|
| 149 |
+
# will require matching the sub-lists of param_shapes for each param group flattened tensor
|
| 150 |
+
|
| 151 |
+
fp32_flat_groups = [
|
| 152 |
+
torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts))
|
| 153 |
+
]
|
| 154 |
+
|
| 155 |
+
return zero_stage, world_size, fp32_flat_groups
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
|
| 159 |
+
"""
|
| 160 |
+
Returns fp32 state_dict reconstructed from ds checkpoint
|
| 161 |
+
|
| 162 |
+
Args:
|
| 163 |
+
- ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
|
| 164 |
+
|
| 165 |
+
"""
|
| 166 |
+
print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
|
| 167 |
+
|
| 168 |
+
optim_files = get_optim_files(ds_checkpoint_dir)
|
| 169 |
+
zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
|
| 170 |
+
print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
|
| 171 |
+
|
| 172 |
+
model_file = get_model_state_file(ds_checkpoint_dir, zero_stage)
|
| 173 |
+
buffers, param_shapes, shared_params, ds_version = parse_model_state(model_file)
|
| 174 |
+
print(f'Parsing checkpoint created by deepspeed=={ds_version}')
|
| 175 |
+
|
| 176 |
+
if zero_stage == 2:
|
| 177 |
+
return _get_fp32_state_dict_from_zero2_checkpoint(world_size, param_shapes, fp32_flat_groups, buffers,
|
| 178 |
+
shared_params)
|
| 179 |
+
elif zero_stage == 3:
|
| 180 |
+
return _get_fp32_state_dict_from_zero3_checkpoint(world_size, param_shapes, fp32_flat_groups, buffers,
|
| 181 |
+
shared_params)
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
def _get_fp32_state_dict_from_zero2_checkpoint(world_size, param_shapes, fp32_flat_groups, buffers, shared_params):
|
| 185 |
+
|
| 186 |
+
# Reconstruction protocol:
|
| 187 |
+
#
|
| 188 |
+
# XXX: document this
|
| 189 |
+
|
| 190 |
+
if debug:
|
| 191 |
+
for i in range(world_size):
|
| 192 |
+
for j in range(len(fp32_flat_groups[0])):
|
| 193 |
+
print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
|
| 194 |
+
|
| 195 |
+
# XXX: memory usage doubles here (zero2)
|
| 196 |
+
num_param_groups = len(fp32_flat_groups[0])
|
| 197 |
+
merged_single_partition_of_fp32_groups = []
|
| 198 |
+
for i in range(num_param_groups):
|
| 199 |
+
merged_partitions = [sd[i] for sd in fp32_flat_groups]
|
| 200 |
+
full_single_fp32_vector = torch.cat(merged_partitions, 0)
|
| 201 |
+
merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
|
| 202 |
+
avail_numel = sum(
|
| 203 |
+
[full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
|
| 204 |
+
|
| 205 |
+
if debug:
|
| 206 |
+
wanted_params = sum([len(shapes) for shapes in param_shapes])
|
| 207 |
+
wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
|
| 208 |
+
# not asserting if there is a mismatch due to possible padding
|
| 209 |
+
print(f"Have {avail_numel} numels to process.")
|
| 210 |
+
print(f"Need {wanted_numel} numels in {wanted_params} params.")
|
| 211 |
+
|
| 212 |
+
state_dict = OrderedDict()
|
| 213 |
+
|
| 214 |
+
# buffers
|
| 215 |
+
state_dict.update(buffers)
|
| 216 |
+
if debug:
|
| 217 |
+
print(f"added {len(buffers)} buffers")
|
| 218 |
+
|
| 219 |
+
# params
|
| 220 |
+
# XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
|
| 221 |
+
# out-of-core computing solution
|
| 222 |
+
total_numel = 0
|
| 223 |
+
total_params = 0
|
| 224 |
+
for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
|
| 225 |
+
offset = 0
|
| 226 |
+
avail_numel = full_single_fp32_vector.numel()
|
| 227 |
+
for name, shape in shapes.items():
|
| 228 |
+
|
| 229 |
+
unpartitioned_numel = shape.numel()
|
| 230 |
+
total_numel += unpartitioned_numel
|
| 231 |
+
total_params += 1
|
| 232 |
+
|
| 233 |
+
if debug:
|
| 234 |
+
print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
|
| 235 |
+
state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
|
| 236 |
+
offset += unpartitioned_numel
|
| 237 |
+
|
| 238 |
+
# Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
|
| 239 |
+
# avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
|
| 240 |
+
# paddings performed in the code it's almost impossible to predict the exact numbers w/o the
|
| 241 |
+
# live optimizer object, so we are checking that the numbers are within the right range
|
| 242 |
+
align_to = 2 * world_size
|
| 243 |
+
|
| 244 |
+
def zero2_align(x):
|
| 245 |
+
return align_to * math.ceil(x / align_to)
|
| 246 |
+
|
| 247 |
+
if debug:
|
| 248 |
+
print(f"original offset={offset}, avail_numel={avail_numel}")
|
| 249 |
+
|
| 250 |
+
offset = zero2_align(offset)
|
| 251 |
+
avail_numel = zero2_align(avail_numel)
|
| 252 |
+
|
| 253 |
+
if debug:
|
| 254 |
+
print(f"aligned offset={offset}, avail_numel={avail_numel}")
|
| 255 |
+
|
| 256 |
+
# Sanity check
|
| 257 |
+
if offset != avail_numel:
|
| 258 |
+
raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
|
| 259 |
+
|
| 260 |
+
# recover shared parameters
|
| 261 |
+
for pair in shared_params:
|
| 262 |
+
state_dict[pair[0]] = state_dict[pair[1]]
|
| 263 |
+
|
| 264 |
+
print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
|
| 265 |
+
|
| 266 |
+
return state_dict
|
| 267 |
+
|
| 268 |
+
|
| 269 |
+
def zero3_partitioned_param_info(unpartitioned_numel, world_size):
|
| 270 |
+
remainder = unpartitioned_numel % world_size
|
| 271 |
+
padding_numel = (world_size - remainder) if remainder else 0
|
| 272 |
+
partitioned_numel = math.ceil(unpartitioned_numel / world_size)
|
| 273 |
+
return partitioned_numel, padding_numel
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def _get_fp32_state_dict_from_zero3_checkpoint(world_size, param_shapes, fp32_flat_groups, buffers, shared_params):
|
| 277 |
+
|
| 278 |
+
# Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
|
| 279 |
+
# param, re-consolidating each param, while dealing with padding if any
|
| 280 |
+
|
| 281 |
+
avail_numel = fp32_flat_groups[0].numel() * world_size
|
| 282 |
+
# merge list of dicts, preserving order
|
| 283 |
+
param_shapes = {k: v for d in param_shapes for k, v in d.items()}
|
| 284 |
+
|
| 285 |
+
if debug:
|
| 286 |
+
for i in range(world_size):
|
| 287 |
+
print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
|
| 288 |
+
|
| 289 |
+
wanted_params = len(param_shapes)
|
| 290 |
+
wanted_numel = sum(shape.numel() for shape in param_shapes.values())
|
| 291 |
+
# not asserting if there is a mismatch due to possible padding
|
| 292 |
+
print(f"Have {avail_numel} numels to process.")
|
| 293 |
+
print(f"Need {wanted_numel} numels in {wanted_params} params.")
|
| 294 |
+
|
| 295 |
+
state_dict = OrderedDict()
|
| 296 |
+
|
| 297 |
+
# buffers
|
| 298 |
+
state_dict.update(buffers)
|
| 299 |
+
if debug:
|
| 300 |
+
print(f"added {len(buffers)} buffers")
|
| 301 |
+
|
| 302 |
+
# params
|
| 303 |
+
# XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
|
| 304 |
+
# out-of-core computing solution
|
| 305 |
+
offset = 0
|
| 306 |
+
total_numel = 0
|
| 307 |
+
total_params = 0
|
| 308 |
+
for name, shape in param_shapes.items():
|
| 309 |
+
|
| 310 |
+
unpartitioned_numel = shape.numel()
|
| 311 |
+
total_numel += unpartitioned_numel
|
| 312 |
+
total_params += 1
|
| 313 |
+
|
| 314 |
+
partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
|
| 315 |
+
|
| 316 |
+
if debug:
|
| 317 |
+
print(
|
| 318 |
+
f"{total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
|
| 319 |
+
)
|
| 320 |
+
|
| 321 |
+
# XXX: memory usage doubles here
|
| 322 |
+
state_dict[name] = torch.cat(
|
| 323 |
+
tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
|
| 324 |
+
0).narrow(0, 0, unpartitioned_numel).view(shape)
|
| 325 |
+
offset += partitioned_numel
|
| 326 |
+
|
| 327 |
+
offset *= world_size
|
| 328 |
+
|
| 329 |
+
# Sanity check
|
| 330 |
+
if offset != avail_numel:
|
| 331 |
+
raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
|
| 332 |
+
|
| 333 |
+
# recover shared parameters
|
| 334 |
+
for pair in shared_params:
|
| 335 |
+
state_dict[pair[0]] = state_dict[pair[1]]
|
| 336 |
+
|
| 337 |
+
print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
|
| 338 |
+
|
| 339 |
+
return state_dict
|
| 340 |
+
|
| 341 |
+
|
| 342 |
+
def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
|
| 343 |
+
"""
|
| 344 |
+
Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
|
| 345 |
+
``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
|
| 346 |
+
via a model hub.
|
| 347 |
+
|
| 348 |
+
Args:
|
| 349 |
+
- ``checkpoint_dir``: path to the desired checkpoint folder
|
| 350 |
+
- ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
|
| 351 |
+
|
| 352 |
+
Returns:
|
| 353 |
+
- pytorch ``state_dict``
|
| 354 |
+
|
| 355 |
+
Note: this approach may not work if your application doesn't have sufficient free CPU memory and
|
| 356 |
+
you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
|
| 357 |
+
the checkpoint.
|
| 358 |
+
|
| 359 |
+
A typical usage might be ::
|
| 360 |
+
|
| 361 |
+
from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
|
| 362 |
+
# do the training and checkpoint saving
|
| 363 |
+
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
|
| 364 |
+
model = model.cpu() # move to cpu
|
| 365 |
+
model.load_state_dict(state_dict)
|
| 366 |
+
# submit to model hub or save the model to share with others
|
| 367 |
+
|
| 368 |
+
In this example the ``model`` will no longer be usable in the deepspeed context of the same
|
| 369 |
+
application. i.e. you will need to re-initialize the deepspeed engine, since
|
| 370 |
+
``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
|
| 371 |
+
|
| 372 |
+
If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
|
| 373 |
+
|
| 374 |
+
"""
|
| 375 |
+
if tag is None:
|
| 376 |
+
latest_path = os.path.join(checkpoint_dir, 'latest')
|
| 377 |
+
if os.path.isfile(latest_path):
|
| 378 |
+
with open(latest_path, 'r') as fd:
|
| 379 |
+
tag = fd.read().strip()
|
| 380 |
+
else:
|
| 381 |
+
raise ValueError(f"Unable to find 'latest' file at {latest_path}")
|
| 382 |
+
|
| 383 |
+
ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
|
| 384 |
+
|
| 385 |
+
if not os.path.isdir(ds_checkpoint_dir):
|
| 386 |
+
raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
|
| 387 |
+
|
| 388 |
+
return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
|
| 389 |
+
|
| 390 |
+
|
| 391 |
+
def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
|
| 392 |
+
"""
|
| 393 |
+
Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
|
| 394 |
+
loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
|
| 395 |
+
|
| 396 |
+
Args:
|
| 397 |
+
- ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
|
| 398 |
+
- ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
|
| 399 |
+
- ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
|
| 400 |
+
"""
|
| 401 |
+
|
| 402 |
+
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
|
| 403 |
+
print(f"Saving fp32 state dict to {output_file}")
|
| 404 |
+
torch.save(state_dict, output_file)
|
| 405 |
+
|
| 406 |
+
|
| 407 |
+
def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
|
| 408 |
+
"""
|
| 409 |
+
1. Put the provided model to cpu
|
| 410 |
+
2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
|
| 411 |
+
3. Load it into the provided model
|
| 412 |
+
|
| 413 |
+
Args:
|
| 414 |
+
- ``model``: the model object to update
|
| 415 |
+
- ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
|
| 416 |
+
- ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
|
| 417 |
+
|
| 418 |
+
Returns:
|
| 419 |
+
- ``model`: modified model
|
| 420 |
+
|
| 421 |
+
Make sure you have plenty of CPU memory available before you call this function. If you don't
|
| 422 |
+
have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
|
| 423 |
+
conveniently placed for you in the checkpoint folder.
|
| 424 |
+
|
| 425 |
+
A typical usage might be ::
|
| 426 |
+
|
| 427 |
+
from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
|
| 428 |
+
model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
|
| 429 |
+
# submit to model hub or save the model to share with others
|
| 430 |
+
|
| 431 |
+
Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
|
| 432 |
+
of the same application. i.e. you will need to re-initialize the deepspeed engine, since
|
| 433 |
+
``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
|
| 434 |
+
|
| 435 |
+
"""
|
| 436 |
+
logger.info(f"Extracting fp32 weights")
|
| 437 |
+
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
|
| 438 |
+
|
| 439 |
+
logger.info(f"Overwriting model with fp32 weights")
|
| 440 |
+
model = model.cpu()
|
| 441 |
+
model.load_state_dict(state_dict, strict=False)
|
| 442 |
+
|
| 443 |
+
return model
|
| 444 |
+
|
| 445 |
+
|
| 446 |
+
if __name__ == "__main__":
|
| 447 |
+
|
| 448 |
+
parser = argparse.ArgumentParser()
|
| 449 |
+
parser.add_argument("checkpoint_dir",
|
| 450 |
+
type=str,
|
| 451 |
+
help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
|
| 452 |
+
parser.add_argument(
|
| 453 |
+
"output_file",
|
| 454 |
+
type=str,
|
| 455 |
+
help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)")
|
| 456 |
+
parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
|
| 457 |
+
args = parser.parse_args()
|
| 458 |
+
|
| 459 |
+
debug = args.debug
|
| 460 |
+
|
| 461 |
+
convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)
|
checkpoint-500/latest
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
pytorch_model
|
checkpoint-500/pytorch_model/mp_rank_00_model_states.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:38c2a608ea3284fd896fba1b0a3c1c3819fb13ab092a66a6b8dbdbcea27b52cb
|
| 3 |
+
size 1719248603
|
checkpoint-500/pytorch_model/zero_pp_rank_0_mp_rank_00_optim_states.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:1f3a0dd08762ba3a2117c1fed194ae6605e040607caf5561b04ec7b8b7d1de47
|
| 3 |
+
size 10314315422
|
checkpoint-500/random_states_0.pkl
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b8a9c619fd8a4e53f5b20dd087d415b859fe434e1f188d04dcaeb245651d8ab4
|
| 3 |
+
size 14599
|
checkpoint-500/scheduler.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a5366cf02c05c445e8083720cac6bb801ec551d5ad3d4561228e99fa34a4b466
|
| 3 |
+
size 563
|
checkpoint-500/unet/config.json
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_class_name": "UNet2DConditionModel",
|
| 3 |
+
"_diffusers_version": "0.16.0.dev0",
|
| 4 |
+
"_name_or_path": "runwayml/stable-diffusion-v1-5",
|
| 5 |
+
"act_fn": "silu",
|
| 6 |
+
"attention_head_dim": 8,
|
| 7 |
+
"block_out_channels": [
|
| 8 |
+
320,
|
| 9 |
+
640,
|
| 10 |
+
1280,
|
| 11 |
+
1280
|
| 12 |
+
],
|
| 13 |
+
"center_input_sample": false,
|
| 14 |
+
"class_embed_type": null,
|
| 15 |
+
"class_embeddings_concat": false,
|
| 16 |
+
"conv_in_kernel": 3,
|
| 17 |
+
"conv_out_kernel": 3,
|
| 18 |
+
"cross_attention_dim": 768,
|
| 19 |
+
"cross_attention_norm": null,
|
| 20 |
+
"down_block_types": [
|
| 21 |
+
"CrossAttnDownBlock2D",
|
| 22 |
+
"CrossAttnDownBlock2D",
|
| 23 |
+
"CrossAttnDownBlock2D",
|
| 24 |
+
"DownBlock2D"
|
| 25 |
+
],
|
| 26 |
+
"downsample_padding": 1,
|
| 27 |
+
"dual_cross_attention": false,
|
| 28 |
+
"encoder_hid_dim": null,
|
| 29 |
+
"flip_sin_to_cos": true,
|
| 30 |
+
"freq_shift": 0,
|
| 31 |
+
"in_channels": 4,
|
| 32 |
+
"layers_per_block": 2,
|
| 33 |
+
"mid_block_only_cross_attention": null,
|
| 34 |
+
"mid_block_scale_factor": 1,
|
| 35 |
+
"mid_block_type": "UNetMidBlock2DCrossAttn",
|
| 36 |
+
"norm_eps": 1e-05,
|
| 37 |
+
"norm_num_groups": 32,
|
| 38 |
+
"num_class_embeds": null,
|
| 39 |
+
"only_cross_attention": false,
|
| 40 |
+
"out_channels": 4,
|
| 41 |
+
"projection_class_embeddings_input_dim": null,
|
| 42 |
+
"resnet_out_scale_factor": 1.0,
|
| 43 |
+
"resnet_skip_time_act": false,
|
| 44 |
+
"resnet_time_scale_shift": "default",
|
| 45 |
+
"sample_size": 64,
|
| 46 |
+
"time_cond_proj_dim": null,
|
| 47 |
+
"time_embedding_act_fn": null,
|
| 48 |
+
"time_embedding_type": "positional",
|
| 49 |
+
"timestep_post_act": null,
|
| 50 |
+
"up_block_types": [
|
| 51 |
+
"UpBlock2D",
|
| 52 |
+
"CrossAttnUpBlock2D",
|
| 53 |
+
"CrossAttnUpBlock2D",
|
| 54 |
+
"CrossAttnUpBlock2D"
|
| 55 |
+
],
|
| 56 |
+
"upcast_attention": false,
|
| 57 |
+
"use_linear_projection": false
|
| 58 |
+
}
|
checkpoint-500/unet/diffusion_pytorch_model.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b9bf7e8a45997ea77a26ad78d42189ace0ec0ce44ad33e0b6449f3fb0021fbe6
|
| 3 |
+
size 1719188507
|
checkpoint-500/zero_to_fp32.py
ADDED
|
@@ -0,0 +1,461 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python
|
| 2 |
+
|
| 3 |
+
# Copyright (c) Microsoft Corporation.
|
| 4 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 5 |
+
|
| 6 |
+
# DeepSpeed Team
|
| 7 |
+
|
| 8 |
+
# This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
|
| 9 |
+
# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
|
| 10 |
+
# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
|
| 11 |
+
# application.
|
| 12 |
+
#
|
| 13 |
+
# example: python zero_to_fp32.py . pytorch_model.bin
|
| 14 |
+
|
| 15 |
+
import argparse
|
| 16 |
+
import torch
|
| 17 |
+
import glob
|
| 18 |
+
import math
|
| 19 |
+
import os
|
| 20 |
+
import re
|
| 21 |
+
from collections import OrderedDict
|
| 22 |
+
|
| 23 |
+
# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
|
| 24 |
+
# DeepSpeed data structures it has to be available in the current python environment.
|
| 25 |
+
from deepspeed.utils import logger
|
| 26 |
+
from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
|
| 27 |
+
FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES)
|
| 28 |
+
|
| 29 |
+
debug = 0
|
| 30 |
+
|
| 31 |
+
# load to cpu
|
| 32 |
+
device = torch.device('cpu')
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def atoi(text):
|
| 36 |
+
return int(text) if text.isdigit() else text
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def natural_keys(text):
|
| 40 |
+
'''
|
| 41 |
+
alist.sort(key=natural_keys) sorts in human order
|
| 42 |
+
http://nedbatchelder.com/blog/200712/human_sorting.html
|
| 43 |
+
(See Toothy's implementation in the comments)
|
| 44 |
+
'''
|
| 45 |
+
return [atoi(c) for c in re.split(r'(\d+)', text)]
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def get_model_state_file(checkpoint_dir, zero_stage):
|
| 49 |
+
if not os.path.isdir(checkpoint_dir):
|
| 50 |
+
raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
|
| 51 |
+
|
| 52 |
+
# there should be only one file
|
| 53 |
+
if zero_stage == 2:
|
| 54 |
+
file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
|
| 55 |
+
elif zero_stage == 3:
|
| 56 |
+
file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
|
| 57 |
+
|
| 58 |
+
if not os.path.exists(file):
|
| 59 |
+
raise FileNotFoundError(f"can't find model states file at '{file}'")
|
| 60 |
+
|
| 61 |
+
return file
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def get_optim_files(checkpoint_dir):
|
| 65 |
+
# XXX: need to test that this simple glob rule works for multi-node setup too
|
| 66 |
+
optim_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*_optim_states.pt")), key=natural_keys)
|
| 67 |
+
|
| 68 |
+
if len(optim_files) == 0:
|
| 69 |
+
raise FileNotFoundError(f"can't find '*_optim_states.pt' files in directory '{checkpoint_dir}'")
|
| 70 |
+
|
| 71 |
+
return optim_files
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def parse_model_state(file):
|
| 75 |
+
state_dict = torch.load(file, map_location=device)
|
| 76 |
+
|
| 77 |
+
if BUFFER_NAMES not in state_dict:
|
| 78 |
+
raise ValueError(f"{file} is not a model state checkpoint")
|
| 79 |
+
buffer_names = state_dict[BUFFER_NAMES]
|
| 80 |
+
if debug:
|
| 81 |
+
print("Found buffers:", buffer_names)
|
| 82 |
+
|
| 83 |
+
# recover just the buffers while restoring them to fp32 if they were saved in fp16
|
| 84 |
+
buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
|
| 85 |
+
param_shapes = state_dict[PARAM_SHAPES]
|
| 86 |
+
|
| 87 |
+
# collect parameters that are included in param_shapes
|
| 88 |
+
param_names = []
|
| 89 |
+
for s in param_shapes:
|
| 90 |
+
for name in s.keys():
|
| 91 |
+
param_names.append(name)
|
| 92 |
+
|
| 93 |
+
# record shared parameters so that they can be recovered based on partners
|
| 94 |
+
# this is because such parameters holding reference only are not saved by optimizer
|
| 95 |
+
shared_params = []
|
| 96 |
+
for param in state_dict["module"]:
|
| 97 |
+
if param not in [*param_names, *buffer_names]:
|
| 98 |
+
for share_param in state_dict["module"]:
|
| 99 |
+
if (state_dict["module"][share_param].data_ptr() == state_dict["module"][param].data_ptr()
|
| 100 |
+
and share_param != param):
|
| 101 |
+
shared_params.append([param, share_param])
|
| 102 |
+
break
|
| 103 |
+
|
| 104 |
+
ds_version = state_dict.get(DS_VERSION, None)
|
| 105 |
+
|
| 106 |
+
return buffers, param_shapes, shared_params, ds_version
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def parse_optim_states(files, ds_checkpoint_dir):
|
| 110 |
+
|
| 111 |
+
total_files = len(files)
|
| 112 |
+
state_dicts = []
|
| 113 |
+
for f in files:
|
| 114 |
+
state_dicts.append(torch.load(f, map_location=device))
|
| 115 |
+
|
| 116 |
+
if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
|
| 117 |
+
raise ValueError(f"{files[0]} is not a zero checkpoint")
|
| 118 |
+
zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
|
| 119 |
+
world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
|
| 120 |
+
|
| 121 |
+
# For ZeRO-2 each param group can have different partition_count as data parallelism for expert
|
| 122 |
+
# parameters can be different from data parallelism for non-expert parameters. So we can just
|
| 123 |
+
# use the max of the partition_count to get the dp world_size.
|
| 124 |
+
|
| 125 |
+
if type(world_size) is list:
|
| 126 |
+
world_size = max(world_size)
|
| 127 |
+
|
| 128 |
+
if world_size != total_files:
|
| 129 |
+
raise ValueError(
|
| 130 |
+
f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
|
| 131 |
+
"Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
|
| 132 |
+
)
|
| 133 |
+
|
| 134 |
+
# the groups are named differently in each stage
|
| 135 |
+
if zero_stage == 2:
|
| 136 |
+
fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
|
| 137 |
+
elif zero_stage == 3:
|
| 138 |
+
fp32_groups_key = FP32_FLAT_GROUPS
|
| 139 |
+
else:
|
| 140 |
+
raise ValueError(f"unknown zero stage {zero_stage}")
|
| 141 |
+
|
| 142 |
+
if zero_stage == 2:
|
| 143 |
+
fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
|
| 144 |
+
elif zero_stage == 3:
|
| 145 |
+
# if there is more than one param group, there will be multiple flattened tensors - one
|
| 146 |
+
# flattened tensor per group - for simplicity merge them into a single tensor
|
| 147 |
+
#
|
| 148 |
+
# XXX: could make the script more memory efficient for when there are multiple groups - it
|
| 149 |
+
# will require matching the sub-lists of param_shapes for each param group flattened tensor
|
| 150 |
+
|
| 151 |
+
fp32_flat_groups = [
|
| 152 |
+
torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts))
|
| 153 |
+
]
|
| 154 |
+
|
| 155 |
+
return zero_stage, world_size, fp32_flat_groups
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
|
| 159 |
+
"""
|
| 160 |
+
Returns fp32 state_dict reconstructed from ds checkpoint
|
| 161 |
+
|
| 162 |
+
Args:
|
| 163 |
+
- ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
|
| 164 |
+
|
| 165 |
+
"""
|
| 166 |
+
print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
|
| 167 |
+
|
| 168 |
+
optim_files = get_optim_files(ds_checkpoint_dir)
|
| 169 |
+
zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
|
| 170 |
+
print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
|
| 171 |
+
|
| 172 |
+
model_file = get_model_state_file(ds_checkpoint_dir, zero_stage)
|
| 173 |
+
buffers, param_shapes, shared_params, ds_version = parse_model_state(model_file)
|
| 174 |
+
print(f'Parsing checkpoint created by deepspeed=={ds_version}')
|
| 175 |
+
|
| 176 |
+
if zero_stage == 2:
|
| 177 |
+
return _get_fp32_state_dict_from_zero2_checkpoint(world_size, param_shapes, fp32_flat_groups, buffers,
|
| 178 |
+
shared_params)
|
| 179 |
+
elif zero_stage == 3:
|
| 180 |
+
return _get_fp32_state_dict_from_zero3_checkpoint(world_size, param_shapes, fp32_flat_groups, buffers,
|
| 181 |
+
shared_params)
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
def _get_fp32_state_dict_from_zero2_checkpoint(world_size, param_shapes, fp32_flat_groups, buffers, shared_params):
|
| 185 |
+
|
| 186 |
+
# Reconstruction protocol:
|
| 187 |
+
#
|
| 188 |
+
# XXX: document this
|
| 189 |
+
|
| 190 |
+
if debug:
|
| 191 |
+
for i in range(world_size):
|
| 192 |
+
for j in range(len(fp32_flat_groups[0])):
|
| 193 |
+
print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
|
| 194 |
+
|
| 195 |
+
# XXX: memory usage doubles here (zero2)
|
| 196 |
+
num_param_groups = len(fp32_flat_groups[0])
|
| 197 |
+
merged_single_partition_of_fp32_groups = []
|
| 198 |
+
for i in range(num_param_groups):
|
| 199 |
+
merged_partitions = [sd[i] for sd in fp32_flat_groups]
|
| 200 |
+
full_single_fp32_vector = torch.cat(merged_partitions, 0)
|
| 201 |
+
merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
|
| 202 |
+
avail_numel = sum(
|
| 203 |
+
[full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
|
| 204 |
+
|
| 205 |
+
if debug:
|
| 206 |
+
wanted_params = sum([len(shapes) for shapes in param_shapes])
|
| 207 |
+
wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
|
| 208 |
+
# not asserting if there is a mismatch due to possible padding
|
| 209 |
+
print(f"Have {avail_numel} numels to process.")
|
| 210 |
+
print(f"Need {wanted_numel} numels in {wanted_params} params.")
|
| 211 |
+
|
| 212 |
+
state_dict = OrderedDict()
|
| 213 |
+
|
| 214 |
+
# buffers
|
| 215 |
+
state_dict.update(buffers)
|
| 216 |
+
if debug:
|
| 217 |
+
print(f"added {len(buffers)} buffers")
|
| 218 |
+
|
| 219 |
+
# params
|
| 220 |
+
# XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
|
| 221 |
+
# out-of-core computing solution
|
| 222 |
+
total_numel = 0
|
| 223 |
+
total_params = 0
|
| 224 |
+
for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
|
| 225 |
+
offset = 0
|
| 226 |
+
avail_numel = full_single_fp32_vector.numel()
|
| 227 |
+
for name, shape in shapes.items():
|
| 228 |
+
|
| 229 |
+
unpartitioned_numel = shape.numel()
|
| 230 |
+
total_numel += unpartitioned_numel
|
| 231 |
+
total_params += 1
|
| 232 |
+
|
| 233 |
+
if debug:
|
| 234 |
+
print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
|
| 235 |
+
state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
|
| 236 |
+
offset += unpartitioned_numel
|
| 237 |
+
|
| 238 |
+
# Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
|
| 239 |
+
# avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
|
| 240 |
+
# paddings performed in the code it's almost impossible to predict the exact numbers w/o the
|
| 241 |
+
# live optimizer object, so we are checking that the numbers are within the right range
|
| 242 |
+
align_to = 2 * world_size
|
| 243 |
+
|
| 244 |
+
def zero2_align(x):
|
| 245 |
+
return align_to * math.ceil(x / align_to)
|
| 246 |
+
|
| 247 |
+
if debug:
|
| 248 |
+
print(f"original offset={offset}, avail_numel={avail_numel}")
|
| 249 |
+
|
| 250 |
+
offset = zero2_align(offset)
|
| 251 |
+
avail_numel = zero2_align(avail_numel)
|
| 252 |
+
|
| 253 |
+
if debug:
|
| 254 |
+
print(f"aligned offset={offset}, avail_numel={avail_numel}")
|
| 255 |
+
|
| 256 |
+
# Sanity check
|
| 257 |
+
if offset != avail_numel:
|
| 258 |
+
raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
|
| 259 |
+
|
| 260 |
+
# recover shared parameters
|
| 261 |
+
for pair in shared_params:
|
| 262 |
+
state_dict[pair[0]] = state_dict[pair[1]]
|
| 263 |
+
|
| 264 |
+
print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
|
| 265 |
+
|
| 266 |
+
return state_dict
|
| 267 |
+
|
| 268 |
+
|
| 269 |
+
def zero3_partitioned_param_info(unpartitioned_numel, world_size):
|
| 270 |
+
remainder = unpartitioned_numel % world_size
|
| 271 |
+
padding_numel = (world_size - remainder) if remainder else 0
|
| 272 |
+
partitioned_numel = math.ceil(unpartitioned_numel / world_size)
|
| 273 |
+
return partitioned_numel, padding_numel
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def _get_fp32_state_dict_from_zero3_checkpoint(world_size, param_shapes, fp32_flat_groups, buffers, shared_params):
|
| 277 |
+
|
| 278 |
+
# Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
|
| 279 |
+
# param, re-consolidating each param, while dealing with padding if any
|
| 280 |
+
|
| 281 |
+
avail_numel = fp32_flat_groups[0].numel() * world_size
|
| 282 |
+
# merge list of dicts, preserving order
|
| 283 |
+
param_shapes = {k: v for d in param_shapes for k, v in d.items()}
|
| 284 |
+
|
| 285 |
+
if debug:
|
| 286 |
+
for i in range(world_size):
|
| 287 |
+
print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
|
| 288 |
+
|
| 289 |
+
wanted_params = len(param_shapes)
|
| 290 |
+
wanted_numel = sum(shape.numel() for shape in param_shapes.values())
|
| 291 |
+
# not asserting if there is a mismatch due to possible padding
|
| 292 |
+
print(f"Have {avail_numel} numels to process.")
|
| 293 |
+
print(f"Need {wanted_numel} numels in {wanted_params} params.")
|
| 294 |
+
|
| 295 |
+
state_dict = OrderedDict()
|
| 296 |
+
|
| 297 |
+
# buffers
|
| 298 |
+
state_dict.update(buffers)
|
| 299 |
+
if debug:
|
| 300 |
+
print(f"added {len(buffers)} buffers")
|
| 301 |
+
|
| 302 |
+
# params
|
| 303 |
+
# XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
|
| 304 |
+
# out-of-core computing solution
|
| 305 |
+
offset = 0
|
| 306 |
+
total_numel = 0
|
| 307 |
+
total_params = 0
|
| 308 |
+
for name, shape in param_shapes.items():
|
| 309 |
+
|
| 310 |
+
unpartitioned_numel = shape.numel()
|
| 311 |
+
total_numel += unpartitioned_numel
|
| 312 |
+
total_params += 1
|
| 313 |
+
|
| 314 |
+
partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
|
| 315 |
+
|
| 316 |
+
if debug:
|
| 317 |
+
print(
|
| 318 |
+
f"{total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
|
| 319 |
+
)
|
| 320 |
+
|
| 321 |
+
# XXX: memory usage doubles here
|
| 322 |
+
state_dict[name] = torch.cat(
|
| 323 |
+
tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
|
| 324 |
+
0).narrow(0, 0, unpartitioned_numel).view(shape)
|
| 325 |
+
offset += partitioned_numel
|
| 326 |
+
|
| 327 |
+
offset *= world_size
|
| 328 |
+
|
| 329 |
+
# Sanity check
|
| 330 |
+
if offset != avail_numel:
|
| 331 |
+
raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
|
| 332 |
+
|
| 333 |
+
# recover shared parameters
|
| 334 |
+
for pair in shared_params:
|
| 335 |
+
state_dict[pair[0]] = state_dict[pair[1]]
|
| 336 |
+
|
| 337 |
+
print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
|
| 338 |
+
|
| 339 |
+
return state_dict
|
| 340 |
+
|
| 341 |
+
|
| 342 |
+
def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
|
| 343 |
+
"""
|
| 344 |
+
Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
|
| 345 |
+
``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
|
| 346 |
+
via a model hub.
|
| 347 |
+
|
| 348 |
+
Args:
|
| 349 |
+
- ``checkpoint_dir``: path to the desired checkpoint folder
|
| 350 |
+
- ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
|
| 351 |
+
|
| 352 |
+
Returns:
|
| 353 |
+
- pytorch ``state_dict``
|
| 354 |
+
|
| 355 |
+
Note: this approach may not work if your application doesn't have sufficient free CPU memory and
|
| 356 |
+
you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
|
| 357 |
+
the checkpoint.
|
| 358 |
+
|
| 359 |
+
A typical usage might be ::
|
| 360 |
+
|
| 361 |
+
from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
|
| 362 |
+
# do the training and checkpoint saving
|
| 363 |
+
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
|
| 364 |
+
model = model.cpu() # move to cpu
|
| 365 |
+
model.load_state_dict(state_dict)
|
| 366 |
+
# submit to model hub or save the model to share with others
|
| 367 |
+
|
| 368 |
+
In this example the ``model`` will no longer be usable in the deepspeed context of the same
|
| 369 |
+
application. i.e. you will need to re-initialize the deepspeed engine, since
|
| 370 |
+
``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
|
| 371 |
+
|
| 372 |
+
If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
|
| 373 |
+
|
| 374 |
+
"""
|
| 375 |
+
if tag is None:
|
| 376 |
+
latest_path = os.path.join(checkpoint_dir, 'latest')
|
| 377 |
+
if os.path.isfile(latest_path):
|
| 378 |
+
with open(latest_path, 'r') as fd:
|
| 379 |
+
tag = fd.read().strip()
|
| 380 |
+
else:
|
| 381 |
+
raise ValueError(f"Unable to find 'latest' file at {latest_path}")
|
| 382 |
+
|
| 383 |
+
ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
|
| 384 |
+
|
| 385 |
+
if not os.path.isdir(ds_checkpoint_dir):
|
| 386 |
+
raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
|
| 387 |
+
|
| 388 |
+
return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
|
| 389 |
+
|
| 390 |
+
|
| 391 |
+
def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
|
| 392 |
+
"""
|
| 393 |
+
Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
|
| 394 |
+
loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
|
| 395 |
+
|
| 396 |
+
Args:
|
| 397 |
+
- ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
|
| 398 |
+
- ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
|
| 399 |
+
- ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
|
| 400 |
+
"""
|
| 401 |
+
|
| 402 |
+
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
|
| 403 |
+
print(f"Saving fp32 state dict to {output_file}")
|
| 404 |
+
torch.save(state_dict, output_file)
|
| 405 |
+
|
| 406 |
+
|
| 407 |
+
def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
|
| 408 |
+
"""
|
| 409 |
+
1. Put the provided model to cpu
|
| 410 |
+
2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
|
| 411 |
+
3. Load it into the provided model
|
| 412 |
+
|
| 413 |
+
Args:
|
| 414 |
+
- ``model``: the model object to update
|
| 415 |
+
- ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
|
| 416 |
+
- ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
|
| 417 |
+
|
| 418 |
+
Returns:
|
| 419 |
+
- ``model`: modified model
|
| 420 |
+
|
| 421 |
+
Make sure you have plenty of CPU memory available before you call this function. If you don't
|
| 422 |
+
have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
|
| 423 |
+
conveniently placed for you in the checkpoint folder.
|
| 424 |
+
|
| 425 |
+
A typical usage might be ::
|
| 426 |
+
|
| 427 |
+
from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
|
| 428 |
+
model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
|
| 429 |
+
# submit to model hub or save the model to share with others
|
| 430 |
+
|
| 431 |
+
Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
|
| 432 |
+
of the same application. i.e. you will need to re-initialize the deepspeed engine, since
|
| 433 |
+
``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
|
| 434 |
+
|
| 435 |
+
"""
|
| 436 |
+
logger.info(f"Extracting fp32 weights")
|
| 437 |
+
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
|
| 438 |
+
|
| 439 |
+
logger.info(f"Overwriting model with fp32 weights")
|
| 440 |
+
model = model.cpu()
|
| 441 |
+
model.load_state_dict(state_dict, strict=False)
|
| 442 |
+
|
| 443 |
+
return model
|
| 444 |
+
|
| 445 |
+
|
| 446 |
+
if __name__ == "__main__":
|
| 447 |
+
|
| 448 |
+
parser = argparse.ArgumentParser()
|
| 449 |
+
parser.add_argument("checkpoint_dir",
|
| 450 |
+
type=str,
|
| 451 |
+
help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
|
| 452 |
+
parser.add_argument(
|
| 453 |
+
"output_file",
|
| 454 |
+
type=str,
|
| 455 |
+
help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)")
|
| 456 |
+
parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
|
| 457 |
+
args = parser.parse_args()
|
| 458 |
+
|
| 459 |
+
debug = args.debug
|
| 460 |
+
|
| 461 |
+
convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)
|
feature_extractor/preprocessor_config.json
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"crop_size": {
|
| 3 |
+
"height": 224,
|
| 4 |
+
"width": 224
|
| 5 |
+
},
|
| 6 |
+
"do_center_crop": true,
|
| 7 |
+
"do_convert_rgb": true,
|
| 8 |
+
"do_normalize": true,
|
| 9 |
+
"do_rescale": true,
|
| 10 |
+
"do_resize": true,
|
| 11 |
+
"feature_extractor_type": "CLIPFeatureExtractor",
|
| 12 |
+
"image_mean": [
|
| 13 |
+
0.48145466,
|
| 14 |
+
0.4578275,
|
| 15 |
+
0.40821073
|
| 16 |
+
],
|
| 17 |
+
"image_processor_type": "CLIPFeatureExtractor",
|
| 18 |
+
"image_std": [
|
| 19 |
+
0.26862954,
|
| 20 |
+
0.26130258,
|
| 21 |
+
0.27577711
|
| 22 |
+
],
|
| 23 |
+
"resample": 3,
|
| 24 |
+
"rescale_factor": 0.00392156862745098,
|
| 25 |
+
"size": {
|
| 26 |
+
"shortest_edge": 224
|
| 27 |
+
}
|
| 28 |
+
}
|
model_index.json
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_class_name": "StableDiffusionPipeline",
|
| 3 |
+
"_diffusers_version": "0.16.0.dev0",
|
| 4 |
+
"feature_extractor": [
|
| 5 |
+
"transformers",
|
| 6 |
+
"CLIPFeatureExtractor"
|
| 7 |
+
],
|
| 8 |
+
"requires_safety_checker": true,
|
| 9 |
+
"safety_checker": [
|
| 10 |
+
"stable_diffusion",
|
| 11 |
+
"StableDiffusionSafetyChecker"
|
| 12 |
+
],
|
| 13 |
+
"scheduler": [
|
| 14 |
+
"diffusers",
|
| 15 |
+
"PNDMScheduler"
|
| 16 |
+
],
|
| 17 |
+
"text_encoder": [
|
| 18 |
+
"transformers",
|
| 19 |
+
"CLIPTextModel"
|
| 20 |
+
],
|
| 21 |
+
"tokenizer": [
|
| 22 |
+
"transformers",
|
| 23 |
+
"CLIPTokenizer"
|
| 24 |
+
],
|
| 25 |
+
"unet": [
|
| 26 |
+
"diffusers",
|
| 27 |
+
"UNet2DConditionModel"
|
| 28 |
+
],
|
| 29 |
+
"vae": [
|
| 30 |
+
"diffusers",
|
| 31 |
+
"AutoencoderKL"
|
| 32 |
+
]
|
| 33 |
+
}
|
safety_checker/config.json
ADDED
|
@@ -0,0 +1,168 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_commit_hash": "39593d5650112b4cc580433f6b0435385882d819",
|
| 3 |
+
"_name_or_path": "/home/iskandre/.cache/huggingface/hub/models--runwayml--stable-diffusion-v1-5/snapshots/39593d5650112b4cc580433f6b0435385882d819/safety_checker",
|
| 4 |
+
"architectures": [
|
| 5 |
+
"StableDiffusionSafetyChecker"
|
| 6 |
+
],
|
| 7 |
+
"initializer_factor": 1.0,
|
| 8 |
+
"logit_scale_init_value": 2.6592,
|
| 9 |
+
"model_type": "clip",
|
| 10 |
+
"projection_dim": 768,
|
| 11 |
+
"text_config": {
|
| 12 |
+
"_name_or_path": "",
|
| 13 |
+
"add_cross_attention": false,
|
| 14 |
+
"architectures": null,
|
| 15 |
+
"attention_dropout": 0.0,
|
| 16 |
+
"bad_words_ids": null,
|
| 17 |
+
"begin_suppress_tokens": null,
|
| 18 |
+
"bos_token_id": 0,
|
| 19 |
+
"chunk_size_feed_forward": 0,
|
| 20 |
+
"cross_attention_hidden_size": null,
|
| 21 |
+
"decoder_start_token_id": null,
|
| 22 |
+
"diversity_penalty": 0.0,
|
| 23 |
+
"do_sample": false,
|
| 24 |
+
"dropout": 0.0,
|
| 25 |
+
"early_stopping": false,
|
| 26 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 27 |
+
"eos_token_id": 2,
|
| 28 |
+
"exponential_decay_length_penalty": null,
|
| 29 |
+
"finetuning_task": null,
|
| 30 |
+
"forced_bos_token_id": null,
|
| 31 |
+
"forced_eos_token_id": null,
|
| 32 |
+
"hidden_act": "quick_gelu",
|
| 33 |
+
"hidden_size": 768,
|
| 34 |
+
"id2label": {
|
| 35 |
+
"0": "LABEL_0",
|
| 36 |
+
"1": "LABEL_1"
|
| 37 |
+
},
|
| 38 |
+
"initializer_factor": 1.0,
|
| 39 |
+
"initializer_range": 0.02,
|
| 40 |
+
"intermediate_size": 3072,
|
| 41 |
+
"is_decoder": false,
|
| 42 |
+
"is_encoder_decoder": false,
|
| 43 |
+
"label2id": {
|
| 44 |
+
"LABEL_0": 0,
|
| 45 |
+
"LABEL_1": 1
|
| 46 |
+
},
|
| 47 |
+
"layer_norm_eps": 1e-05,
|
| 48 |
+
"length_penalty": 1.0,
|
| 49 |
+
"max_length": 20,
|
| 50 |
+
"max_position_embeddings": 77,
|
| 51 |
+
"min_length": 0,
|
| 52 |
+
"model_type": "clip_text_model",
|
| 53 |
+
"no_repeat_ngram_size": 0,
|
| 54 |
+
"num_attention_heads": 12,
|
| 55 |
+
"num_beam_groups": 1,
|
| 56 |
+
"num_beams": 1,
|
| 57 |
+
"num_hidden_layers": 12,
|
| 58 |
+
"num_return_sequences": 1,
|
| 59 |
+
"output_attentions": false,
|
| 60 |
+
"output_hidden_states": false,
|
| 61 |
+
"output_scores": false,
|
| 62 |
+
"pad_token_id": 1,
|
| 63 |
+
"prefix": null,
|
| 64 |
+
"problem_type": null,
|
| 65 |
+
"projection_dim": 512,
|
| 66 |
+
"pruned_heads": {},
|
| 67 |
+
"remove_invalid_values": false,
|
| 68 |
+
"repetition_penalty": 1.0,
|
| 69 |
+
"return_dict": true,
|
| 70 |
+
"return_dict_in_generate": false,
|
| 71 |
+
"sep_token_id": null,
|
| 72 |
+
"suppress_tokens": null,
|
| 73 |
+
"task_specific_params": null,
|
| 74 |
+
"temperature": 1.0,
|
| 75 |
+
"tf_legacy_loss": false,
|
| 76 |
+
"tie_encoder_decoder": false,
|
| 77 |
+
"tie_word_embeddings": true,
|
| 78 |
+
"tokenizer_class": null,
|
| 79 |
+
"top_k": 50,
|
| 80 |
+
"top_p": 1.0,
|
| 81 |
+
"torch_dtype": null,
|
| 82 |
+
"torchscript": false,
|
| 83 |
+
"transformers_version": "4.28.1",
|
| 84 |
+
"typical_p": 1.0,
|
| 85 |
+
"use_bfloat16": false,
|
| 86 |
+
"vocab_size": 49408
|
| 87 |
+
},
|
| 88 |
+
"torch_dtype": "float32",
|
| 89 |
+
"transformers_version": null,
|
| 90 |
+
"vision_config": {
|
| 91 |
+
"_name_or_path": "",
|
| 92 |
+
"add_cross_attention": false,
|
| 93 |
+
"architectures": null,
|
| 94 |
+
"attention_dropout": 0.0,
|
| 95 |
+
"bad_words_ids": null,
|
| 96 |
+
"begin_suppress_tokens": null,
|
| 97 |
+
"bos_token_id": null,
|
| 98 |
+
"chunk_size_feed_forward": 0,
|
| 99 |
+
"cross_attention_hidden_size": null,
|
| 100 |
+
"decoder_start_token_id": null,
|
| 101 |
+
"diversity_penalty": 0.0,
|
| 102 |
+
"do_sample": false,
|
| 103 |
+
"dropout": 0.0,
|
| 104 |
+
"early_stopping": false,
|
| 105 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 106 |
+
"eos_token_id": null,
|
| 107 |
+
"exponential_decay_length_penalty": null,
|
| 108 |
+
"finetuning_task": null,
|
| 109 |
+
"forced_bos_token_id": null,
|
| 110 |
+
"forced_eos_token_id": null,
|
| 111 |
+
"hidden_act": "quick_gelu",
|
| 112 |
+
"hidden_size": 1024,
|
| 113 |
+
"id2label": {
|
| 114 |
+
"0": "LABEL_0",
|
| 115 |
+
"1": "LABEL_1"
|
| 116 |
+
},
|
| 117 |
+
"image_size": 224,
|
| 118 |
+
"initializer_factor": 1.0,
|
| 119 |
+
"initializer_range": 0.02,
|
| 120 |
+
"intermediate_size": 4096,
|
| 121 |
+
"is_decoder": false,
|
| 122 |
+
"is_encoder_decoder": false,
|
| 123 |
+
"label2id": {
|
| 124 |
+
"LABEL_0": 0,
|
| 125 |
+
"LABEL_1": 1
|
| 126 |
+
},
|
| 127 |
+
"layer_norm_eps": 1e-05,
|
| 128 |
+
"length_penalty": 1.0,
|
| 129 |
+
"max_length": 20,
|
| 130 |
+
"min_length": 0,
|
| 131 |
+
"model_type": "clip_vision_model",
|
| 132 |
+
"no_repeat_ngram_size": 0,
|
| 133 |
+
"num_attention_heads": 16,
|
| 134 |
+
"num_beam_groups": 1,
|
| 135 |
+
"num_beams": 1,
|
| 136 |
+
"num_channels": 3,
|
| 137 |
+
"num_hidden_layers": 24,
|
| 138 |
+
"num_return_sequences": 1,
|
| 139 |
+
"output_attentions": false,
|
| 140 |
+
"output_hidden_states": false,
|
| 141 |
+
"output_scores": false,
|
| 142 |
+
"pad_token_id": null,
|
| 143 |
+
"patch_size": 14,
|
| 144 |
+
"prefix": null,
|
| 145 |
+
"problem_type": null,
|
| 146 |
+
"projection_dim": 512,
|
| 147 |
+
"pruned_heads": {},
|
| 148 |
+
"remove_invalid_values": false,
|
| 149 |
+
"repetition_penalty": 1.0,
|
| 150 |
+
"return_dict": true,
|
| 151 |
+
"return_dict_in_generate": false,
|
| 152 |
+
"sep_token_id": null,
|
| 153 |
+
"suppress_tokens": null,
|
| 154 |
+
"task_specific_params": null,
|
| 155 |
+
"temperature": 1.0,
|
| 156 |
+
"tf_legacy_loss": false,
|
| 157 |
+
"tie_encoder_decoder": false,
|
| 158 |
+
"tie_word_embeddings": true,
|
| 159 |
+
"tokenizer_class": null,
|
| 160 |
+
"top_k": 50,
|
| 161 |
+
"top_p": 1.0,
|
| 162 |
+
"torch_dtype": null,
|
| 163 |
+
"torchscript": false,
|
| 164 |
+
"transformers_version": "4.28.1",
|
| 165 |
+
"typical_p": 1.0,
|
| 166 |
+
"use_bfloat16": false
|
| 167 |
+
}
|
| 168 |
+
}
|
safety_checker/pytorch_model.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:16d28f2b37109f222cdc33620fdd262102ac32112be0352a7f77e9614b35a394
|
| 3 |
+
size 1216064769
|
scheduler/scheduler_config.json
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_class_name": "PNDMScheduler",
|
| 3 |
+
"_diffusers_version": "0.16.0.dev0",
|
| 4 |
+
"beta_end": 0.012,
|
| 5 |
+
"beta_schedule": "scaled_linear",
|
| 6 |
+
"beta_start": 0.00085,
|
| 7 |
+
"clip_sample": false,
|
| 8 |
+
"num_train_timesteps": 1000,
|
| 9 |
+
"prediction_type": "epsilon",
|
| 10 |
+
"set_alpha_to_one": false,
|
| 11 |
+
"skip_prk_steps": true,
|
| 12 |
+
"steps_offset": 1,
|
| 13 |
+
"trained_betas": null
|
| 14 |
+
}
|
text_encoder/config.json
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_name_or_path": "runwayml/stable-diffusion-v1-5",
|
| 3 |
+
"architectures": [
|
| 4 |
+
"CLIPTextModel"
|
| 5 |
+
],
|
| 6 |
+
"attention_dropout": 0.0,
|
| 7 |
+
"bos_token_id": 0,
|
| 8 |
+
"dropout": 0.0,
|
| 9 |
+
"eos_token_id": 2,
|
| 10 |
+
"hidden_act": "quick_gelu",
|
| 11 |
+
"hidden_size": 768,
|
| 12 |
+
"initializer_factor": 1.0,
|
| 13 |
+
"initializer_range": 0.02,
|
| 14 |
+
"intermediate_size": 3072,
|
| 15 |
+
"layer_norm_eps": 1e-05,
|
| 16 |
+
"max_position_embeddings": 77,
|
| 17 |
+
"model_type": "clip_text_model",
|
| 18 |
+
"num_attention_heads": 12,
|
| 19 |
+
"num_hidden_layers": 12,
|
| 20 |
+
"pad_token_id": 1,
|
| 21 |
+
"projection_dim": 768,
|
| 22 |
+
"torch_dtype": "float16",
|
| 23 |
+
"transformers_version": "4.28.1",
|
| 24 |
+
"vocab_size": 49408
|
| 25 |
+
}
|
text_encoder/pytorch_model.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:6a34f30098988d85dc0fb0fc272a842ebcf552e2ebc6ce4adbcf3695d08e8a90
|
| 3 |
+
size 246188833
|
tokenizer/merges.txt
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
tokenizer/special_tokens_map.json
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token": {
|
| 3 |
+
"content": "<|startoftext|>",
|
| 4 |
+
"lstrip": false,
|
| 5 |
+
"normalized": true,
|
| 6 |
+
"rstrip": false,
|
| 7 |
+
"single_word": false
|
| 8 |
+
},
|
| 9 |
+
"eos_token": {
|
| 10 |
+
"content": "<|endoftext|>",
|
| 11 |
+
"lstrip": false,
|
| 12 |
+
"normalized": true,
|
| 13 |
+
"rstrip": false,
|
| 14 |
+
"single_word": false
|
| 15 |
+
},
|
| 16 |
+
"pad_token": "<|endoftext|>",
|
| 17 |
+
"unk_token": {
|
| 18 |
+
"content": "<|endoftext|>",
|
| 19 |
+
"lstrip": false,
|
| 20 |
+
"normalized": true,
|
| 21 |
+
"rstrip": false,
|
| 22 |
+
"single_word": false
|
| 23 |
+
}
|
| 24 |
+
}
|
tokenizer/tokenizer_config.json
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"add_prefix_space": false,
|
| 3 |
+
"bos_token": {
|
| 4 |
+
"__type": "AddedToken",
|
| 5 |
+
"content": "<|startoftext|>",
|
| 6 |
+
"lstrip": false,
|
| 7 |
+
"normalized": true,
|
| 8 |
+
"rstrip": false,
|
| 9 |
+
"single_word": false
|
| 10 |
+
},
|
| 11 |
+
"clean_up_tokenization_spaces": true,
|
| 12 |
+
"do_lower_case": true,
|
| 13 |
+
"eos_token": {
|
| 14 |
+
"__type": "AddedToken",
|
| 15 |
+
"content": "<|endoftext|>",
|
| 16 |
+
"lstrip": false,
|
| 17 |
+
"normalized": true,
|
| 18 |
+
"rstrip": false,
|
| 19 |
+
"single_word": false
|
| 20 |
+
},
|
| 21 |
+
"errors": "replace",
|
| 22 |
+
"model_max_length": 77,
|
| 23 |
+
"pad_token": "<|endoftext|>",
|
| 24 |
+
"tokenizer_class": "CLIPTokenizer",
|
| 25 |
+
"unk_token": {
|
| 26 |
+
"__type": "AddedToken",
|
| 27 |
+
"content": "<|endoftext|>",
|
| 28 |
+
"lstrip": false,
|
| 29 |
+
"normalized": true,
|
| 30 |
+
"rstrip": false,
|
| 31 |
+
"single_word": false
|
| 32 |
+
}
|
| 33 |
+
}
|
tokenizer/vocab.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
unet/config.json
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_class_name": "UNet2DConditionModel",
|
| 3 |
+
"_diffusers_version": "0.16.0.dev0",
|
| 4 |
+
"_name_or_path": "runwayml/stable-diffusion-v1-5",
|
| 5 |
+
"act_fn": "silu",
|
| 6 |
+
"attention_head_dim": 8,
|
| 7 |
+
"block_out_channels": [
|
| 8 |
+
320,
|
| 9 |
+
640,
|
| 10 |
+
1280,
|
| 11 |
+
1280
|
| 12 |
+
],
|
| 13 |
+
"center_input_sample": false,
|
| 14 |
+
"class_embed_type": null,
|
| 15 |
+
"class_embeddings_concat": false,
|
| 16 |
+
"conv_in_kernel": 3,
|
| 17 |
+
"conv_out_kernel": 3,
|
| 18 |
+
"cross_attention_dim": 768,
|
| 19 |
+
"cross_attention_norm": null,
|
| 20 |
+
"down_block_types": [
|
| 21 |
+
"CrossAttnDownBlock2D",
|
| 22 |
+
"CrossAttnDownBlock2D",
|
| 23 |
+
"CrossAttnDownBlock2D",
|
| 24 |
+
"DownBlock2D"
|
| 25 |
+
],
|
| 26 |
+
"downsample_padding": 1,
|
| 27 |
+
"dual_cross_attention": false,
|
| 28 |
+
"encoder_hid_dim": null,
|
| 29 |
+
"flip_sin_to_cos": true,
|
| 30 |
+
"freq_shift": 0,
|
| 31 |
+
"in_channels": 4,
|
| 32 |
+
"layers_per_block": 2,
|
| 33 |
+
"mid_block_only_cross_attention": null,
|
| 34 |
+
"mid_block_scale_factor": 1,
|
| 35 |
+
"mid_block_type": "UNetMidBlock2DCrossAttn",
|
| 36 |
+
"norm_eps": 1e-05,
|
| 37 |
+
"norm_num_groups": 32,
|
| 38 |
+
"num_class_embeds": null,
|
| 39 |
+
"only_cross_attention": false,
|
| 40 |
+
"out_channels": 4,
|
| 41 |
+
"projection_class_embeddings_input_dim": null,
|
| 42 |
+
"resnet_out_scale_factor": 1.0,
|
| 43 |
+
"resnet_skip_time_act": false,
|
| 44 |
+
"resnet_time_scale_shift": "default",
|
| 45 |
+
"sample_size": 64,
|
| 46 |
+
"time_cond_proj_dim": null,
|
| 47 |
+
"time_embedding_act_fn": null,
|
| 48 |
+
"time_embedding_type": "positional",
|
| 49 |
+
"timestep_post_act": null,
|
| 50 |
+
"up_block_types": [
|
| 51 |
+
"UpBlock2D",
|
| 52 |
+
"CrossAttnUpBlock2D",
|
| 53 |
+
"CrossAttnUpBlock2D",
|
| 54 |
+
"CrossAttnUpBlock2D"
|
| 55 |
+
],
|
| 56 |
+
"upcast_attention": false,
|
| 57 |
+
"use_linear_projection": false
|
| 58 |
+
}
|
unet/diffusion_pytorch_model.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:8ba459ca07ef238e953e658fe514317150271e1df15d70d19faf04f3249cd6d5
|
| 3 |
+
size 1719188507
|
vae/config.json
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_class_name": "AutoencoderKL",
|
| 3 |
+
"_diffusers_version": "0.16.0.dev0",
|
| 4 |
+
"_name_or_path": "/home/iskandre/.cache/huggingface/hub/models--runwayml--stable-diffusion-v1-5/snapshots/39593d5650112b4cc580433f6b0435385882d819/vae",
|
| 5 |
+
"act_fn": "silu",
|
| 6 |
+
"block_out_channels": [
|
| 7 |
+
128,
|
| 8 |
+
256,
|
| 9 |
+
512,
|
| 10 |
+
512
|
| 11 |
+
],
|
| 12 |
+
"down_block_types": [
|
| 13 |
+
"DownEncoderBlock2D",
|
| 14 |
+
"DownEncoderBlock2D",
|
| 15 |
+
"DownEncoderBlock2D",
|
| 16 |
+
"DownEncoderBlock2D"
|
| 17 |
+
],
|
| 18 |
+
"in_channels": 3,
|
| 19 |
+
"latent_channels": 4,
|
| 20 |
+
"layers_per_block": 2,
|
| 21 |
+
"norm_num_groups": 32,
|
| 22 |
+
"out_channels": 3,
|
| 23 |
+
"sample_size": 512,
|
| 24 |
+
"scaling_factor": 0.18215,
|
| 25 |
+
"up_block_types": [
|
| 26 |
+
"UpDecoderBlock2D",
|
| 27 |
+
"UpDecoderBlock2D",
|
| 28 |
+
"UpDecoderBlock2D",
|
| 29 |
+
"UpDecoderBlock2D"
|
| 30 |
+
]
|
| 31 |
+
}
|
vae/diffusion_pytorch_model.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:af27ea858349760ebe3311953e0bfe8d6fd257dc9537ae0b2b938c262132a2c6
|
| 3 |
+
size 334711857
|