Get the latest state dict from a root checkpoint directory. Args: checkpoint_root_path (str): The root checkpoint directory. Returns: Tuple[str, int]: The state dict path and the iteration of the state dict. If the state dict does not exist, return (None, 0).
(
checkpoint_root_path: str,
trainer_type: str = "verl",
)
| 213 | |
| 214 | |
| 215 | def get_latest_state_dict( |
| 216 | checkpoint_root_path: str, |
| 217 | trainer_type: str = "verl", |
| 218 | ) -> Tuple[str, int]: |
| 219 | """Get the latest state dict from a root checkpoint directory. |
| 220 | |
| 221 | Args: |
| 222 | checkpoint_root_path (str): The root checkpoint directory. |
| 223 | |
| 224 | Returns: |
| 225 | Tuple[str, int]: The state dict path and the iteration of the state dict. |
| 226 | If the state dict does not exist, return (None, 0). |
| 227 | """ |
| 228 | if trainer_type != "verl": |
| 229 | raise NotImplementedError(f"Unsupported trainer type {trainer_type}") |
| 230 | latest_state_dict_iteration_path = os.path.join( |
| 231 | checkpoint_root_path, "latest_state_dict_iteration.txt" |
| 232 | ) |
| 233 | if os.path.exists(latest_state_dict_iteration_path): |
| 234 | with open(latest_state_dict_iteration_path, "r", encoding="utf-8") as f: |
| 235 | iteration = f.read().strip() |
| 236 | state_dict_path = os.path.join( |
| 237 | checkpoint_root_path, f"global_step_{iteration}", "actor" |
| 238 | ) |
| 239 | return state_dict_path, int(iteration) |
| 240 | return None, 0 # type: ignore |
| 241 | |
| 242 | |
| 243 | def has_huggingface_model_weights(checkpoint_path: str) -> bool: |