(model,
load_dir,
tag,
load_module_strict=True,
load_optimizer_states=True,
load_lr_scheduler_states=True)
| 72 | |
| 73 | |
| 74 | def _load_checkpoint(model, |
| 75 | load_dir, |
| 76 | tag, |
| 77 | load_module_strict=True, |
| 78 | load_optimizer_states=True, |
| 79 | load_lr_scheduler_states=True): |
| 80 | |
| 81 | load_path = model._get_ckpt_name(load_dir, tag) |
| 82 | |
| 83 | if not os.path.exists(load_path): |
| 84 | return None, None |
| 85 | |
| 86 | checkpoint = torch.load( |
| 87 | load_path, |
| 88 | map_location=lambda storage, loc: storage, |
| 89 | weights_only=True) |
| 90 | |
| 91 | model.load_module_state_dict( |
| 92 | state_dict=checkpoint['module'], strict=load_module_strict) |
| 93 | if not model.zero_optimization() and load_optimizer_states: |
| 94 | if model.fp16_enabled(): |
| 95 | model.optimizer.load_state_dict( |
| 96 | checkpoint['optimizer'], |
| 97 | load_optimizer_states=load_optimizer_states) |
| 98 | elif load_optimizer_states: |
| 99 | model.optimizer.load_state_dict(checkpoint['optimizer']) |
| 100 | |
| 101 | if load_lr_scheduler_states and model.lr_scheduler is not None: |
| 102 | model.lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) |
| 103 | |
| 104 | model.csr_tensor_module_names = checkpoint['csr_tensor_module_names'] |
| 105 | model.global_steps = checkpoint['global_steps'] |
| 106 | model.global_samples = checkpoint.get( |
| 107 | 'global_samples', model.global_steps * model.train_batch_size()) |
| 108 | model.skipped_steps = checkpoint['skipped_steps'] |
| 109 | model.loaded_checkpoint_mp_world_size = checkpoint['mp_world_size'] |
| 110 | model.loaded_checkpoint_dp_world_size = checkpoint['dp_world_size'] |
| 111 | deepspeed_states = [ |
| 112 | 'module', 'optimizer', 'lr_scheduler', 'csr_tensor_module_names', |
| 113 | 'skipped_steps', 'global_steps', 'dp_world_size', 'mp_world_size' |
| 114 | ] |
| 115 | client_state = { |
| 116 | key: value |
| 117 | for key, value in checkpoint.items() if key not in deepspeed_states |
| 118 | } |
| 119 | |
| 120 | return load_path, client_state |
no test coverage detected
searching dependent graphs…