Load a model checkpoint and return the iteration.
(model, optimizer, lr_scheduler, load_arg='load')
| 173 | |
| 174 | |
| 175 | def load_checkpoint(model, optimizer, lr_scheduler, load_arg='load'): |
| 176 | """Load a model checkpoint and return the iteration.""" |
| 177 | args = get_args() |
| 178 | load_dir = getattr(args, load_arg) |
| 179 | |
| 180 | if isinstance(model, torchDDP): |
| 181 | model = model.module |
| 182 | # Read the tracker file and set the iteration. |
| 183 | tracker_filename = get_checkpoint_tracker_filename(load_dir) |
| 184 | |
| 185 | # If no tracker file, return iretation zero. |
| 186 | if not os.path.isfile(tracker_filename): |
| 187 | print_rank_0('WARNING: could not find the metadata file {} '.format( |
| 188 | tracker_filename)) |
| 189 | print_rank_0(' will not load any checkpoints and will start from ' |
| 190 | 'random') |
| 191 | return 0 |
| 192 | |
| 193 | # Otherwise, read the tracker file and either set the iteration or |
| 194 | # mark it as a release checkpoint. |
| 195 | iteration = 0 |
| 196 | release = False |
| 197 | with open(tracker_filename, 'r') as f: |
| 198 | metastring = f.read().strip() |
| 199 | try: |
| 200 | iteration = int(metastring) |
| 201 | except ValueError: |
| 202 | release = metastring == 'release' |
| 203 | if not release: |
| 204 | print_rank_0('ERROR: Invalid metadata file {}. Exiting'.format( |
| 205 | tracker_filename)) |
| 206 | sys.exit() |
| 207 | |
| 208 | assert iteration > 0 or release, 'error parsing metadata file {}'.format( |
| 209 | tracker_filename) |
| 210 | |
| 211 | if args.deepspeed: |
| 212 | checkpoint_name, state_dict = model.load_checkpoint(load_dir) |
| 213 | |
| 214 | if checkpoint_name is None: |
| 215 | if mpu.get_data_parallel_rank() == 0: |
| 216 | print("Unable to load checkpoint.") |
| 217 | return iteration |
| 218 | |
| 219 | else: |
| 220 | # Checkpoint. |
| 221 | checkpoint_name = get_checkpoint_name(load_dir, iteration, release) |
| 222 | if mpu.get_data_parallel_rank() == 0: |
| 223 | print('global rank {} is loading checkpoint {}'.format( |
| 224 | torch.distributed.get_rank(), checkpoint_name)) |
| 225 | |
| 226 | # Load the checkpoint. |
| 227 | try: |
| 228 | state_dict = torch.load(checkpoint_name, map_location='cpu') |
| 229 | except ModuleNotFoundError: |
| 230 | # For backward compatibility. |
| 231 | print_rank_0(' > deserializing using the old code structure ...') |
| 232 | sys.modules['fp16.loss_scaler'] = sys.modules[ |
no test coverage detected