| 161 | |
| 162 | @staticmethod |
| 163 | def save_model(rank, model, out_path: str, dcp=False, lora=False): |
| 164 | |
| 165 | # if lora: |
| 166 | # state_dict = { |
| 167 | # name: param |
| 168 | # for name, param in model.named_parameters() |
| 169 | # if "lora" in name |
| 170 | # } |
| 171 | # else: |
| 172 | state_dict = get_model_state_dict( |
| 173 | model, |
| 174 | options=StateDictOptions( |
| 175 | full_state_dict=not dcp, |
| 176 | cpu_offload=True, |
| 177 | ), |
| 178 | ) |
| 179 | if lora: |
| 180 | state_dict = {key: value for key, value in state_dict.items() if 'lora' in key.lower()} |
| 181 | |
| 182 | if dcp: |
| 183 | DCP.save(state_dict, checkpoint_id=out_path) |
| 184 | elif rank == 0: |
| 185 | torch.save(state_dict, out_path) |
| 186 | del state_dict |
| 187 | if dist.is_initialized(): |
| 188 | dist.barrier() |
| 189 | return |
| 190 | |
| 191 | @staticmethod |
| 192 | def save_optimizer(model, opt, out_path: str, dcp=False): |