(self,
mp_world_size,
mp_rank,
module_key=AUTO_MODULE_KEY,
is_pipe_parallel=False,
quantize=False,
quantize_bits=8,
quantize_groups=64,
mlp_extra_grouping=True)
| 55 | self.check_ckpt_list() |
| 56 | |
| 57 | def load(self, |
| 58 | mp_world_size, |
| 59 | mp_rank, |
| 60 | module_key=AUTO_MODULE_KEY, |
| 61 | is_pipe_parallel=False, |
| 62 | quantize=False, |
| 63 | quantize_bits=8, |
| 64 | quantize_groups=64, |
| 65 | mlp_extra_grouping=True): |
| 66 | self.module_key = module_key |
| 67 | num_ckpt = len(self.ckpt_list) |
| 68 | idx = mp_rank * num_ckpt // mp_world_size |
| 69 | """ We have multiple cases to handle here for both training and inference: |
| 70 | 1. PipeModule loading mp_rank_*.pt files, is_pipe_parallel=True, module_key is not None |
| 71 | a. if no mp_size/pp_size resizing occurs, for both training & inference, loading |
| 72 | the mp_rank related checkpoint directly. |
| 73 | b. if has mp_size/pp_size resizing, only Megatron model inference is supported, |
| 74 | in this case each mp_rank_*.pt have same content, we will load the first checkpoint |
| 75 | file (idx=0), to avoid idx exceeding file list boundary. |
| 76 | |
| 77 | 2. PipeModule loading layer_*.pt files, is_pipe_parallel=True, module_key is None |
| 78 | a. if no mp_size resizing occurs, for both training & inference, loading |
| 79 | the mp_rank related checkpoint directly. |
| 80 | b. if has mp_size resizing, only Megatron model inference is supported, |
| 81 | checkpoint file(s) will be merged/split according to mp_rank, mp_world_size and |
| 82 | checkpoint file list. |
| 83 | |
| 84 | 3. Non-PipeModule loading mp_rank_*.pt files, is_pipe_parallel=False |
| 85 | Same with case (2). |
| 86 | """ |
| 87 | if is_pipe_parallel and module_key is not None and mp_world_size != num_ckpt: |
| 88 | mp_world_size = num_ckpt |
| 89 | idx = 0 |
| 90 | |
| 91 | load_path = self.ckpt_list[idx] |
| 92 | |
| 93 | merge_count = 1 |
| 94 | if num_ckpt == mp_world_size: |
| 95 | assert os.path.exists(load_path) |
| 96 | #logger.info(f'rank: {mp_rank} loading checkpoint: {load_path}') |
| 97 | sd = self.checkpoint_engine.load(load_path, map_location=lambda storage, \ |
| 98 | loc: storage) |
| 99 | |
| 100 | if quantize: |
| 101 | quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) |
| 102 | sd_module, all_scales = quantizer.sd_quantize_megatron(self.get_module(sd), quantize_bits, |
| 103 | quantize_groups) |
| 104 | self.set_module(sd, sd_module) |
| 105 | else: |
| 106 | all_scales = None |
| 107 | elif num_ckpt > mp_world_size: |
| 108 | sd, all_scales, merge_count = self.merge_state_dict(mp_world_size, mp_rank, quantize, \ |
| 109 | quantize_bits, quantize_groups, mlp_extra_grouping) |
| 110 | else: |
| 111 | sd, all_scales = self.split_state_dict(mp_world_size, mp_rank, quantize, quantize_bits, \ |
| 112 | quantize_groups, mlp_extra_grouping) |
| 113 | return load_path, sd, (all_scales, merge_count) |
| 114 |
no test coverage detected