Record the memory state. Args: with_options (str, optional): The options to include in the memory state. Defaults to "". create (bool, optional): Whether to create a new memory record file. Defaults to False. Returns: None
(self, with_options: str = "", create: bool = False)
| 259 | self.point(with_options="params,grads,os_params", create=True) |
| 260 | |
| 261 | def point(self, with_options: str = "", create: bool = False) -> None: |
| 262 | """ |
| 263 | Record the memory state. |
| 264 | |
| 265 | Args: |
| 266 | with_options (str, optional): The options to include in the memory state. Defaults to "". |
| 267 | create (bool, optional): Whether to create a new memory record file. Defaults to False. |
| 268 | |
| 269 | Returns: |
| 270 | None |
| 271 | """ |
| 272 | now = time.time() |
| 273 | file = f"{self._log_folder}/memory.log" |
| 274 | |
| 275 | if with_options == "all": |
| 276 | options = ["params", "grads", "os_params", "os_state", "activation_base"] |
| 277 | else: |
| 278 | options = with_options.split(",") |
| 279 | |
| 280 | total_mem = ( |
| 281 | self._param_mem_state.total_mem |
| 282 | + self._grad_mem_state.total_mem |
| 283 | + self._os_params_mem_state.total_mem |
| 284 | + self._os_state_mem_state.total_mem |
| 285 | + self._activation_mem |
| 286 | ) / mb |
| 287 | |
| 288 | # Generate summary information for memory state |
| 289 | summary_info = ( |
| 290 | f"total_memory: {total_mem:.2f} MB" |
| 291 | + "\n" |
| 292 | + f"params_memory: {self._param_mem_state.total_mem / mb:.2f} MB, " |
| 293 | + f"grads_memory: {self._grad_mem_state.total_mem / mb:.2f} MB, " |
| 294 | + f"os_params_memory: {self._os_params_mem_state.total_mem / mb:.2f} MB, " |
| 295 | + f"os_state_memory: {self._os_state_mem_state.total_mem / mb:.2f} MB, " |
| 296 | + f"activation_memory: {self._activation_mem / mb:.2f} MB" |
| 297 | ) |
| 298 | |
| 299 | # Generate layout information based on selected options |
| 300 | layout_info = "" |
| 301 | if "params" in options: |
| 302 | layout_info += "params_layout:\n" + self._param_mem_state.dump() |
| 303 | if "grads" in options: |
| 304 | layout_info += "grads_layout:\n" + self._grad_mem_state.dump() |
| 305 | if "os_params" in options: |
| 306 | layout_info += "os_params_layout:\n" + self._os_params_mem_state.dump() |
| 307 | if "os_state" in options: |
| 308 | layout_info += "os_state_layout:\n" + self._os_state_mem_state.dump() |
| 309 | if "activation_base" in options: |
| 310 | layout_info += "activation_base_layout:\n" + self._activation_base_mems.dump() |
| 311 | |
| 312 | # Write memory state information to log file |
| 313 | file_mode = "w" if create else "a" |
| 314 | with open(file, file_mode, encoding="utf-8") as writer: |
| 315 | writer.write( |
| 316 | "Memory State:\n" + f"time: {now - self._record_start_time}\n" + "---summary---\n" + summary_info + "\n" |
| 317 | ) |
| 318 | if layout_info != "": |
no test coverage detected