Class SimpleMemoryProfiler

internlm/utils/simple_memory_profiler.py:205–591

A memory profiler for a llm model. Args: model (torch.nn.Module): The model to profile. optimizer (torch.optim.Optimizer): The optimizer used for training the model. log_file (str): The file to write the memory state information to. total_steps: number of st

class SimpleMemoryProfiler:
    """
    A memory profiler for a llm model.

    Args:
        model (torch.nn.Module): The model to profile.
        optimizer (torch.optim.Optimizer): The optimizer used for training the model.
        log_file (str): The file to write the memory state information to.
        total_steps: number of steps to trace.
    """

    def __init__(
        self,
        model: torch.nn.Module,
        optimizer: torch.optim.Optimizer,
        log_folder: str,
        total_steps: int = 5,
    ):
        self._model, self._num_model_chunks = _unpack_naive_wrapper(model)
        self._optimizer = optimizer
        self._log_folder = log_folder
        self._remaining_steps = total_steps

        self._stoped = False
        self._record_start_time = time.time()

        # For activation memory state.

        self._activation_mem: int = 0
        self._activation_mem_max: int = 0
        self._activation_base_mems = ActivationMemState(self._num_model_chunks)

        # Check or create log folder
        os.makedirs(self._log_folder, exist_ok=True)

        # Register activation memory tracking hooks
        if self._num_model_chunks > 1:
            for chunk_id in range(self._num_model_chunks):
                self._register_activation_trace_hooks(chunk_id, self._model[chunk_id])
        else:
            self._register_activation_trace_hooks(0, self._model)

        # Calculate static parameter cuda memory
        self._param_mem_state = SimpleMemState("param_mem")
        self._calc_tensor_memory(self._param_mem_state, self._model.named_parameters())
        # Calculate static grad cuda memory
        self._grad_mem_state = SimpleMemState("grad_mem")
        self._calc_tensor_memory(self._grad_mem_state, self._model.named_parameters(), True)
        # Calculate static optimizer state cuda memory
        self._os_params_mem_state = SimpleMemState("os_params_mem")
        self._os_state_mem_state = SimpleMemState("os_state_mem")
        self._calc_tensor_group_memory(self._os_params_mem_state, list(enumerate(self._optimizer.param_groups)))

        # Generate the first memory record
        self.point(with_options="params,grads,os_params", create=True)

    def point(self, with_options: str = "", create: bool = False) -> None:
        """