Inkdrop/gpt2-property-classifier
029
1#!/usr/bin/env python2 3# This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets4# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in5# the future. Once extracted, the weights don't require DeepSpeed and can be used in any6# application.7#8# example: python zero_to_fp32.py . pytorch_model.bin9 10import argparse11import torch12import glob13import math14import os15from collections import OrderedDict16 17# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with18# DeepSpeed data structures it has to be available in the current python environment.19import deepspeed20from deepspeed.utils import logger21 22debug = 023 24# load to cpu25device = torch.device('cpu')26 27 28def get_model_state_file(checkpoint_dir, zero_stage):29 if not os.path.isdir(checkpoint_dir):30 raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")31 32 # there should be only one file33 if zero_stage == 2:34 file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")35 elif zero_stage == 3:36 file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")37 38 if not os.path.exists(file):39 raise FileNotFoundError(f"can't find model states file at '{file}'")40 41 return file42 43 44def get_optim_files(checkpoint_dir):45 # XXX: need to test that this simple glob rule works for multi-node setup too46 optim_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*_optim_states.pt")))47 48 if len(optim_files) == 0:49 raise FileNotFoundError(50 f"can't find '*_optim_states.pt' files in directory '{checkpoint_dir}'")51 52 return optim_files53 54 55def parse_model_state(file):56 state_dict = torch.load(file, map_location=device)57 58 if "buffer_names" not in state_dict:59 raise ValueError(f"{file} is not a model state checkpoint")60 buffer_names = state_dict["buffer_names"]61 if debug:62 print("Found buffers:", buffer_names)63 64 # recover just the buffers while restoring them to fp32 if they were saved in fp1665 buffers = {66 k: v.float()67 for k,68 v in state_dict["module"].items() if k in buffer_names69 }70 return buffers71 72 73def parse_optim_states(files, ds_checkpoint_dir):74 75 total_files = len(files)76 state_dicts = []77 for f in files:78 state_dicts.append(torch.load(f, map_location=device))79 80 if not "zero_stage" in state_dicts[0]['optimizer_state_dict']:81 raise ValueError(f"{files[0]} is not a zero checkpoint")82 zero_stage = state_dicts[0]['optimizer_state_dict']["zero_stage"]83 world_size = state_dicts[0]['optimizer_state_dict']["partition_count"]84 param_shapes = state_dicts[0]["param_shapes"]85 # For ZeRO-2 each param group can have different partition_count as data parallelism for expert86 # parameters can be different from data parallelism for non-expert parameters. So we can just87 # use the max of the partition_count to get the dp world_size.88 89 if type(world_size) is list:90 world_size = max(world_size)91 92 if world_size != total_files:93 raise ValueError(94 f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "95 "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."96 )97 98 # the groups are named differently in each stage99 if zero_stage == 2:100 fp32_groups_key = "single_partition_of_fp32_groups"101 elif zero_stage == 3:102 fp32_groups_key = "fp32_flat_groups"103 else:104 raise ValueError(f"unknown zero stage {zero_stage}")105 106 if zero_stage == 2:107 fp32_flat_groups = [108 state_dicts[i]['optimizer_state_dict'][fp32_groups_key]109 for i in range(len(state_dicts))110 ]111 elif zero_stage == 3:112 # if there is more than one param group, there will be multiple flattened tensors - one113 # flattened tensor per group - for simplicity merge them into a single tensor114 #115 # XXX: could make the script more memory efficient for when there are multiple groups - it116 # will require matching the sub-lists of param_shapes for each param group flattened tensor117 118 fp32_flat_groups = [119 torch.cat(state_dicts[i]['optimizer_state_dict'][fp32_groups_key],120 0) for i in range(len(state_dicts))121 ]122 123 return zero_stage, world_size, param_shapes, fp32_flat_groups124 125 126def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):127 """128 Returns fp32 state_dict reconstructed from ds checkpoint129 130 Args:131 - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)132 133 """134 print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")135 136 optim_files = get_optim_files(ds_checkpoint_dir)137 zero_stage, world_size, param_shapes, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)138 print(139 f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")140 141 model_file = get_model_state_file(ds_checkpoint_dir, zero_stage)142 buffers = parse_model_state(model_file)143 144 if zero_stage == 2:145 return _get_fp32_state_dict_from_zero2_checkpoint(world_size,146 param_shapes,147 fp32_flat_groups,148 buffers)149 elif zero_stage == 3:150 return _get_fp32_state_dict_from_zero3_checkpoint(world_size,151 param_shapes,152 fp32_flat_groups,153 buffers)154 155 156def _get_fp32_state_dict_from_zero2_checkpoint(world_size,157 param_shapes,158 fp32_flat_groups,159 buffers):160 161 # Reconstruction protocol:162 #163 # XXX: document this164 165 if debug:166 for i in range(world_size):167 for j in range(len(fp32_flat_groups[0])):168 print(f"fp32_flat_groups[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")169 170 # XXX: memory usage doubles here (zero2)171 num_param_groups = len(fp32_flat_groups[0])172 merged_single_partition_of_fp32_groups = []173 for i in range(num_param_groups):174 merged_partitions = [sd[i] for sd in fp32_flat_groups]175 full_single_fp32_vector = torch.cat(merged_partitions, 0)176 merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)177 avail_numel = sum([178 full_single_fp32_vector.numel()179 for full_single_fp32_vector in merged_single_partition_of_fp32_groups180 ])181 182 if debug:183 wanted_params = sum([len(shapes) for shapes in param_shapes])184 wanted_numel = sum(185 [sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])186 # not asserting if there is a mismatch due to possible padding187 print(f"Have {avail_numel} numels to process.")188 print(f"Need {wanted_numel} numels in {wanted_params} params.")189 190 state_dict = OrderedDict()191 192 # buffers193 state_dict.update(buffers)194 if debug:195 print(f"added {len(buffers)} buffers")196 197 # params198 # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support199 # out-of-core computing solution200 total_numel = 0201 total_params = 0202 for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):203 offset = 0204 avail_numel = full_single_fp32_vector.numel()205 for name, shape in shapes.items():206 207 unpartitioned_numel = shape.numel()208 total_numel += unpartitioned_numel209 total_params += 1210 211 if debug:212 print(213 f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} "214 )215 state_dict[name] = full_single_fp32_vector.narrow(216 0,217 offset,218 unpartitioned_numel).view(shape)219 offset += unpartitioned_numel220 221 # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and222 # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex223 # paddings performed in the code it's almost impossible to predict the exact numbers w/o the224 # live optimizer object, so we are checking that the numbers are within the right range225 align_to = 2 * world_size226 227 def zero2_align(x):228 return align_to * math.ceil(x / align_to)229 230 if debug:231 print(f"original offset={offset}, avail_numel={avail_numel}")232 233 offset = zero2_align(offset)234 avail_numel = zero2_align(avail_numel)235 236 if debug:237 print(f"aligned offset={offset}, avail_numel={avail_numel}")238 239 # Sanity check240 if offset != avail_numel:241 raise ValueError(242 f"consumed {offset} numels out of {avail_numel} - something is wrong")243 244 print(245 f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements"246 )247 248 return state_dict249 250 251def zero3_partitioned_param_info(unpartitioned_numel, world_size):252 remainder = unpartitioned_numel % world_size253 padding_numel = (world_size - remainder) if remainder else 0254 partitioned_numel = math.ceil(unpartitioned_numel / world_size)255 return partitioned_numel, padding_numel256 257 258def _get_fp32_state_dict_from_zero3_checkpoint(world_size,259 param_shapes,260 fp32_flat_groups,261 buffers):262 263 # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each264 # param, re-consolidating each param, while dealing with padding if any265 266 avail_numel = fp32_flat_groups[0].numel() * world_size267 # merge list of dicts, preserving order268 param_shapes = {k: v for d in param_shapes for k, v in d.items()}269 270 if debug:271 for i in range(world_size):272 print(f"fp32_flat_groups[{i}].shape={fp32_flat_groups[i].shape}")273 274 wanted_params = len(param_shapes)275 wanted_numel = sum(shape.numel() for shape in param_shapes.values())276 # not asserting if there is a mismatch due to possible padding277 print(f"Have {avail_numel} numels to process.")278 print(f"Need {wanted_numel} numels in {wanted_params} params.")279 280 state_dict = OrderedDict()281 282 # buffers283 state_dict.update(buffers)284 if debug:285 print(f"added {len(buffers)} buffers")286 287 # params288 # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support289 # out-of-core computing solution290 offset = 0291 total_numel = 0292 total_params = 0293 for name, shape in param_shapes.items():294 295 unpartitioned_numel = shape.numel()296 total_numel += unpartitioned_numel297 total_params += 1298 299 partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)300 301 if debug:302 print(303 f"{total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"304 )305 306 # XXX: memory usage doubles here307 state_dict[name] = torch.cat(308 tuple(fp32_flat_groups[i].narrow(0,309 offset,310 partitioned_numel)311 for i in range(world_size)),312 0).narrow(0,313 0,314 unpartitioned_numel).view(shape)315 offset += partitioned_numel316 317 offset *= world_size318 319 # Sanity check320 if offset != avail_numel:321 raise ValueError(322 f"consumed {offset} numels out of {avail_numel} - something is wrong")323 324 print(325 f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements"326 )327 328 return state_dict329 330 331def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):332 """333 Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with334 ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example335 via a model hub.336 337 Args:338 - ``checkpoint_dir``: path to the desired checkpoint folder339 - ``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``340 341 Returns:342 - pytorch ``state_dict``343 344 Note: this approach may not work if your application doesn't have sufficient free CPU memory and345 you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with346 the checkpoint.347 348 A typical usage might be ::349 350 from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint351 # do the training and checkpoint saving352 state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu353 model = model.cpu() # move to cpu354 model.load_state_dict(state_dict)355 # submit to model hub or save the model to share with others356 357 In this example the ``model`` will no longer be usable in the deepspeed context of the same358 application. i.e. you will need to re-initialize the deepspeed engine, since359 ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.360 361 If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.362 363 """364 if tag is None:365 latest_path = os.path.join(checkpoint_dir, 'latest')366 if os.path.isfile(latest_path):367 with open(latest_path, 'r') as fd:368 tag = fd.read().strip()369 else:370 raise ValueError(f"Unable to find 'latest' file at {latest_path}")371 372 ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)373 374 if not os.path.isdir(ds_checkpoint_dir):375 raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")376 377 return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)378 379 380def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):381 """382 Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be383 loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.384 385 Args:386 - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)387 - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)388 - ``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``389 """390 391 state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)392 print(f"Saving fp32 state dict to {output_file}")393 torch.save(state_dict, output_file)394 395 396def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):397 """398 1. Put the provided model to cpu399 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``400 3. Load it into the provided model401 402 Args:403 - ``model``: the model object to update404 - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)405 - ``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``406 407 Returns:408 - ``model`: modified model409 410 Make sure you have plenty of CPU memory available before you call this function. If you don't411 have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it412 conveniently placed for you in the checkpoint folder.413 414 A typical usage might be ::415 416 from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint417 model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)418 # submit to model hub or save the model to share with others419 420 Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context421 of the same application. i.e. you will need to re-initialize the deepspeed engine, since422 ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.423 424 """425 logger.info(f"Extracting fp32 weights")426 state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)427 428 logger.info(f"Overwriting model with fp32 weights")429 model = model.cpu()430 model.load_state_dict(state_dict, strict=False)431 432 return model433 434 435if __name__ == "__main__":436 437 parser = argparse.ArgumentParser()438 parser.add_argument(439 "checkpoint_dir",440 type=str,441 help="path to the desired checkpoint folder, e.g., path/checkpoint-12")442 parser.add_argument(443 "output_file",444 type=str,445 help=446 "path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)"447 )448 parser.add_argument("-d", "--debug", action='store_true', help="enable debug")449 args = parser.parse_args()450 451 debug = args.debug452 453 convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)454