DeLVM / InternLM /internlm /utils /simple_memory_profiler.py

Upload folder using huggingface_hub

ee3e701 verified over 1 year ago

24.7 kB

import os
import time
from collections import OrderedDict
from functools import partial, reduce
from typing import Any, Dict, List, Tuple

import pyecharts
import torch

from internlm.core.naive_amp import NaiveAMPModel

mb = 1024 * 1024


class SimpleMemState:
    """
    A class to represent the memory state of a model layer.
    Args:
        layer_name (str): The name of the layer.
        layer_mem (int): The memory usage of the layer in bytes.
    """

    def __init__(self, layer_name: str, layer_mem: int = 0) -> None:
        self.layer_name = layer_name

        # Memory status of the current model layer.
        self._layer_mem: int = layer_mem
        # Total memory status of the model and sub-models, initialized with layer memory.
        self._total_mem: int = self._layer_mem
        # SimpleMemState of sub-models.
        self.sub_model_stats = OrderedDict()

    @property
    def layer_mem(self) -> int:
        """
        Get the memory usage of the layer.
        Returns:
            int: The memory usage of the layer in bytes.
        """
        return self._layer_mem

    @layer_mem.setter
    def layer_mem(self, new_layer_mem: int) -> None:
        """
        Set the memory usage of the layer.
        Args:
            new_layer_mem (int): The new memory usage of the layer in bytes.
        """
        diff = new_layer_mem - self._layer_mem
        self._layer_mem = new_layer_mem
        self._total_mem += diff

    @property
    def total_mem(self) -> int:
        """
        Get the total memory usage of the model and sub-models.
        Returns:
            int: The total memory usage in bytes.
        """
        return self._total_mem

    def add(self, layer_name: str, layer_mem: int = 0, flush: bool = True) -> None:
        """
        Add a layer to the memory state.
        Args:
            layer_name (str): The name of the layer.
            layer_mem (int, optional): The memory usage of the layer in bytes. Defaults to 0.
            flush (bool, optional): Whether to update the total memory usage. Defaults to True.
        """
        path = layer_name.split(".")

        target = self.find_layer_state(path, create=True)
        target.layer_mem = layer_mem

        if flush:
            self.update_total_memory()

    def delete(self, layer_name: str, flush: bool = True) -> None:
        """
        Delete a layer from the memory state.
        Args:
            layer_name (str): The name of the layer.
            flush (bool, optional): Whether to update the total memory usage. Defaults to True.
        """
        path = layer_name.split(".")
        assert len(path) >= 2, f"Only support deleting non-root layers, layer_name: {layer_name}"

        parent_path = path[0:-1]
        layer = path[-1]
        parent = self.find_layer_state(parent_path)

        if parent is not None and layer in parent.sub_model_stats:
            del parent.sub_model_stats[layer]

        if flush:
            self.update_total_memory()

    def update_total_memory(self) -> None:
        """
        Update the total memory usage of the model and sub-models.
        """
        self._total_mem = self._layer_mem

        for stat in self.sub_model_stats.values():
            # Update sub-model status first.
            stat.update_total_memory()
            # Add sub-model total_mem to model total_mem.
            self._total_mem += stat._total_mem

    def find_layer_state(self, path: Tuple[str], create: bool = False) -> "SimpleMemState":
        """
        Find the memory state of a layer.
        Args:
            path (Tuple[str]): The path to the layer.
            create (bool, optional): Whether to create the layer if it doesn't exist. Defaults to False.
        Returns:
            SimpleMemState: The memory state of the layer.
        """
        current_node = self

        for _node in path:
            if _node not in current_node.sub_model_stats:
                if not create:
                    return None
                # Create a layer node.
                current_node.sub_model_stats[_node] = SimpleMemState(_node)

            current_node = current_node.sub_model_stats[_node]

        return current_node

    def dump(self, prefix: str = "") -> str:
        """
        Dump the memory state of the model and sub-models.
        Args:
            prefix (str, optional): The prefix to add to the layer names. Defaults to "".
        Returns:
            str: The memory state information.
        """
        cur_prefix = prefix + "." + self.layer_name if prefix != "" else self.layer_name
        res = f"layer: {cur_prefix}, layer_mem: {self.layer_mem / mb:.2f} MB, total_mem: {self.total_mem / mb:.2f} MB\n"

        for sub_layer in self.sub_model_stats.values():
            res += sub_layer.dump(cur_prefix)

        return res

    def to_json(self, base: int = 1024 * 1024) -> dict:
        """
        Convert the memory state to a JSON structure.
        Returns:
            dict: The JSON structure of the memory state.
        """
        children = [child.to_json() for child in self.sub_model_stats.values()]
        if len(children) == 0:
            return {"name": self.layer_name, "value": self.layer_mem // base}
        else:
            return {"name": self.layer_name, "children": children}


class ActivationMemState:
    """
    Activation Memory State
    """

    def __init__(self, num_chunks: int) -> None:
        self._num_chunks = num_chunks

        self.inited: List[bool] = [False for _ in range(num_chunks)]
        self.states: List[SimpleMemState] = [SimpleMemState(f"activations_{idx}") for idx in range(num_chunks)]

    @property
    def total_mem(self) -> int:
        return sum(state.total_mem for state in self.states)

    def dump(self, prefix: str = "") -> str:
        return reduce(lambda x, y: x + y, [state.dump(prefix) for state in self.states])

    def to_json(self, base: int = 1024 * 1024) -> List:
        return [state.to_json(base) for state in self.states]


def _unpack_naive_wrapper(model: torch.nn.Module) -> Tuple[torch.nn.Module, int]:
    num_chunks = len(model) if isinstance(model, torch.nn.ModuleList) else 1

    if num_chunks > 1:
        model = torch.nn.ModuleList([_model.model if isinstance(_model, NaiveAMPModel) else _model for _model in model])
    else:
        model = model.model if isinstance(model, NaiveAMPModel) else model

    return model, num_chunks


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:
        """
        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
        """
        now = time.time()
        file = f"{self._log_folder}/memory.log"

        if with_options == "all":
            options = ["params", "grads", "os_params", "os_state", "activation_base"]
        else:
            options = with_options.split(",")

        total_mem = (
            self._param_mem_state.total_mem
            + self._grad_mem_state.total_mem
            + self._os_params_mem_state.total_mem
            + self._os_state_mem_state.total_mem
            + self._activation_mem
        ) / mb

        # Generate summary information for memory state
        summary_info = (
            f"total_memory: {total_mem:.2f} MB"
            + "\n"
            + f"params_memory: {self._param_mem_state.total_mem / mb:.2f} MB, "
            + f"grads_memory: {self._grad_mem_state.total_mem / mb:.2f} MB, "
            + f"os_params_memory: {self._os_params_mem_state.total_mem / mb:.2f} MB, "
            + f"os_state_memory: {self._os_state_mem_state.total_mem / mb:.2f} MB, "
            + f"activation_memory: {self._activation_mem / mb:.2f} MB"
        )

        # Generate layout information based on selected options
        layout_info = ""
        if "params" in options:
            layout_info += "params_layout:\n" + self._param_mem_state.dump()
        if "grads" in options:
            layout_info += "grads_layout:\n" + self._grad_mem_state.dump()
        if "os_params" in options:
            layout_info += "os_params_layout:\n" + self._os_params_mem_state.dump()
        if "os_state" in options:
            layout_info += "os_state_layout:\n" + self._os_state_mem_state.dump()
        if "activation_base" in options:
            layout_info += "activation_base_layout:\n" + self._activation_base_mems.dump()

        # Write memory state information to log file
        file_mode = "w" if create else "a"
        with open(file, file_mode, encoding="utf-8") as writer:
            writer.write(
                "Memory State:\n" + f"time: {now - self._record_start_time}\n" + "---summary---\n" + summary_info + "\n"
            )
            if layout_info != "":
                writer.write("---Layout---\n" + layout_info)
            writer.write("\n")

    def step(self) -> None:
        """
        Update the memory state of the optimizer state.
        Returns:
            None
        """
        if self._stoped:
            return

        self._remaining_steps -= 1
        if self._remaining_steps == 0:
            self._stoped = True

        # Update os state memory usage
        self._os_state_mem_state = SimpleMemState("os_state_mem")
        self._calc_tensor_group_memory(self._os_state_mem_state, list(self._optimizer.state_dict()["state"].items()))

        if not self._stoped:
            # Do we need to print os_state_layout every time? Is it always constant?
            self.point(with_options="os_state")
        else:
            # Dump memory layout
            self.point(with_options="all")
            # Generate sunburst charts
            self._render_sunburst_chart(self._param_mem_state.to_json()["children"], "params_memory_sunburst")
            self._render_sunburst_chart(self._grad_mem_state.to_json()["children"], "grads_memory_sunburst")
            self._render_sunburst_chart(
                [self._os_params_mem_state.to_json(), self._os_state_mem_state.to_json()],
                "os_memory_sunburst",
            )
            self._render_sunburst_chart(self._activation_base_mems.to_json(), "activation_memory_sunburst")
            # Generate summary sunburst chart
            summary_sunburst_data = [
                {"name": "params", "value": self._param_mem_state.total_mem // mb},
                {"name": "grads", "value": self._grad_mem_state.total_mem // mb},
                {"name": "os_params", "value": self._os_params_mem_state.total_mem // mb},
                {"name": "os_state", "value": self._os_state_mem_state.total_mem // mb},
                {"name": "activation", "value": self._activation_mem_max // mb},
            ]

            self._render_sunburst_chart(summary_sunburst_data, "summary_sunburst")

    def _render_sunburst_chart(self, data: Any, name: str) -> None:
        pyecharts.charts.Sunburst(init_opts=pyecharts.options.InitOpts(width="1000px", height="1000px")).add(
            name,
            data_pair=data,
            highlight_policy="ancestor",
            radius=[0, "95%"],
            levels=[
                {},
                {
                    "r0": "10%",
                    "r": "35%",
                    "itemStyle": {"borderWidth": 3},
                    "label": {"align": "left"},
                },
                {"r0": "35%", "r": "55%", "label": {"align": "left"}},
                {"r0": "55%", "r": "70%", "label": {"align": "left"}},
                {"r0": "70%", "r": "80%", "label": {"align": "left"}},
                {"r0": "80%", "r": "90%", "label": {"align": "left"}},
                {
                    "r0": "90%",
                    "r": "92%",
                    "label": {"position": "outside", "padding": 3, "silent": False},
                    "itemStyle": {"borderWidth": 3},
                },
            ],
        ).set_global_opts(title_opts=pyecharts.options.TitleOpts(title="CUDA Memory")).set_series_opts(
            label_opts=pyecharts.options.LabelOpts(formatter="{b}")
        ).render(
            f"{self._log_folder}/{name}.html"
        )

    def _inner_activation_trace_hook(
        self,
        chunk_id: int,
        layer_name: str,
        model: Any,
        inputs: Any,
        output: torch.Tensor,
    ) -> None:
        """
        Hook function to trace the activation memory usage for a inner layer.
        Args:
            layer_name (str): The name of the layer.
            model (Any): The model.
            inputs (Any): The inputs to the layer.
            output (torch.Tensor): The output tensor.
        Returns:
            None
        """
        del model, inputs
        assert isinstance(output, torch.Tensor), f"Invalid output type: {type(output)}"

        if self._stoped or self._activation_base_mems.inited[chunk_id]:
            return

        # Delay updating the total_mem of activation_base_mem here, it will be handled in the forward ending hook.
        self._activation_base_mems.states[chunk_id].add(
            layer_name, output.element_size() * output.nelement(), flush=False
        )

    def _activation_trace_hook_forward(self, chunk_id: int, model: Any, inputs: Any, output: torch.Tensor) -> None:
        """
        Hook function to trace the activation memory usage for a forward pass.
        Args:
            model (Any): The model.
            inputs (Any): The inputs to the model.
            output (torch.Tensor): The output tensor.
        Returns:
            None
        """
        del model, inputs
        assert isinstance(output, torch.Tensor), f"invalid output type: {type(output)}"

        if self._stoped:
            return

        # Check if the activation memory has been initialized
        if self._activation_base_mems.inited[chunk_id] is False:
            self._activation_base_mems.inited[chunk_id] = True
            # Update the total memory of the activation base memory state
            self._activation_base_mems.states[chunk_id].update_total_memory()
            # Set with_options to "activation_base" to include activation_base_layout in the memory dump
            with_options = "activation_base"
        else:
            with_options = ""

        # Accumulate activation memory usage for each forward pass
        self._activation_mem += self._activation_base_mems.states[chunk_id].total_mem
        if self._activation_mem > self._activation_mem_max:
            self._activation_mem_max = self._activation_mem

        # Trigger a memory record
        self.point(with_options)

    def _activation_tarce_hook_backward(self, chunk_id: int, model: Any, inputs: Any, grad_outputs: Any) -> None:
        """
        Hook function to trace the activation memory usage for a backward pass.
        Args:
            model (Any): The model.
            inputs (Any): The inputs to the model.
            grad_outputs (Any): The gradients of the outputs.
        Returns:
            None
        """
        del model, inputs, grad_outputs

        if self._stoped:
            return

        # Release activation memory usage for each backward pass
        self._activation_mem -= self._activation_base_mems.states[chunk_id].total_mem

        # Trigger a memory record
        self.point()

    def _register_activation_trace_hooks(self, chunk_id: int, model_chunk: torch.nn.Module) -> None:
        """
        Register activation trace hooks for the model and each submodule in the model.
        """

        # Register inner activation trace hooks for each submodule in the model
        for layer_name, sub_model in model_chunk.named_modules():
            # Register the hook
            if len(sub_model._modules) != 0:
                continue  # TODO: in some special cases, we may need some additional configuration to correct

            sub_model.register_forward_hook(partial(self._inner_activation_trace_hook, chunk_id, layer_name))

        # Register a forward hook for the main model to track activation memory usage
        model_chunk.register_forward_hook(partial(self._activation_trace_hook_forward, chunk_id))
        # Register a backward hook for the main model to release activation memory usage
        model_chunk.register_full_backward_hook(partial(self._activation_tarce_hook_backward, chunk_id))

    def _calc_tensor_memory(
        self, root_stat: SimpleMemState, named_tensors: Dict[str, torch.Tensor], require_grad: bool = False
    ) -> None:
        """
        Calculate the memory usage of tensors and update the memory state.
        Args:
            root_stat (SimpleMemState): The root memory state.
            named_tensors (Dict[str, torch.Tensor]): A dictionary containing the named tensors.
            require_grad (bool, optional): Whether to consider tensors with gradients. Defaults to False.
        Returns:
            None
        """
        for name, tensor in named_tensors:
            if require_grad and not tensor.requires_grad:
                continue

            layer_splits = name.split(sep=".")
            layer_stat = root_stat.find_layer_state(layer_splits, create=True)
            layer_stat.layer_mem = tensor.element_size() * tensor.nelement()

        root_stat.update_total_memory()

    def _calc_tensor_group_memory(self, root_stat: SimpleMemState, tensor_groups: List[Tuple[int, torch.Tensor]]):
        """
        Calculate the memory usage of a group of tensors.
        Args:
            root_stat (SimpleMemState): The root memory state.
            tensor_groups (List[Tuple[int, torch.Tensor]]): A list of tuples containing the tensor groups.
        Returns:
            None
        """

        def _normalize_helper(named_tensors: Dict[str, Any]) -> List[Tuple[str, Any]]:
            """
            Normalize the named tensors.
            Args:
                named_tensors (Dict[str, Any]): The named tensors to normalize.
            Returns:
                List[Tuple[str, Any]]: The normalized named tensors.
            """
            res = {}

            for name, tensors in named_tensors.items():
                if isinstance(tensors, torch.Tensor):
                    res[name] = tensors
                elif isinstance(tensors, (list, tuple)):
                    for index, tensor in enumerate(tensors):
                        res[f"{name}.{index}"] = tensor
                elif isinstance(tensors, dict):
                    for subname, tensor in tensors.items():
                        res[f"{name}.{subname}"] = tensor
                else:
                    raise TypeError(f"unsupported normalize value type: {type(tensors)}")

            return list(res.items())

        def _value_check(tensor_or_tensors):
            """
            Check if the input is a tensor or a collection of tensors.
            Args:
                tensor_or_tensors (Any): The input to check.
            Returns:
                bool: True if the input is a tensor or a collection of tensors, False otherwise.
            """
            if torch.is_tensor(tensor_or_tensors):
                return True
            elif isinstance(tensor_or_tensors, (list, tuple)) and all(torch.is_tensor(x) for x in tensor_or_tensors):
                return True
            elif isinstance(tensor_or_tensors, dict) and all(torch.is_tensor(x) for x in tensor_or_tensors.values()):
                return True
            else:
                return False

        # Calculate the memory usage of a group of tensors.
        for idx, tensors in tensor_groups:
            # Normalize the named tensors
            named_tensors = {f"{idx}.{k}": v for k, v in tensors.items() if _value_check(v)}
            named_tensors = _normalize_helper(named_tensors)
            # Calculate the memory usage of the tensors and update the memory state
            self._calc_tensor_memory(root_stat, named_tensors)


if __name__ == "__main__":

    class SimpleModel(torch.nn.Module):
        """
        A simple model with three linear layers.
        Args:
            skip_layer2 (bool, optional): Whether to skip layer2. Defaults to False.
        """

        def __init__(self, skip_layer2: bool = False):
            super().__init__()
            self.layer1 = torch.nn.Linear(5120, 5120, True)
            self.layer3 = torch.nn.Linear(5120, 5120, False)

            if skip_layer2:
                self.layer2 = None
            else:
                self.layer2 = SimpleModel(skip_layer2=True)

        def forward(self, inputs: torch.Tensor) -> torch.Tensor:
            """
            Forward pass of the model.
            Args:
                inputs (torch.Tensor): The input tensor.
            Returns:
                torch.Tensor: The output tensor.
            """
            output1 = self.layer1(inputs)
            if self.layer2 is not None:
                output2 = self.layer2(output1)
            else:
                output2 = output1
            output = self.layer3(output2)

            return output

    def _simple_schedule(_num_chunks, _model_chunks, _input) -> torch.Tensor:
        if _num_chunks > 1:
            _output = _input
            for _model_chunk in _model_chunks:
                _output = _model_chunk(_output)
        else:
            _output = _model_chunks(_input)

        return _output

    # num_chunks config
    _num_chunks = 1

    # init model and optimizer
    if _num_chunks > 1:
        _chunks = [SimpleModel(skip_layer2=idx % 2 == 0) for idx in range(_num_chunks)]
        _model = torch.nn.ModuleList(_chunks).cuda()
    else:
        _model: torch.nn.Module = SimpleModel().cuda()
    _optimizer = torch.optim.Adam(_model.parameters())

    # init profiler
    profiler = SimpleMemoryProfiler(_model, _optimizer, "./test_simple_memory_profiler", total_steps=1)

    _optimizer.zero_grad()

    # inputs
    x1 = torch.randn((128, 5120)).cuda()
    x2 = torch.randn((128, 5120)).cuda()
    # forward
    out1 = _simple_schedule(_num_chunks, _model, x1)
    out2 = _simple_schedule(_num_chunks, _model, x2)
    # backward
    out1.mean().backward()
    out2.mean().backward()

    _optimizer.step()

    # Update the optimizer state memory usage and record the memory state
    profiler.step()

	import os
	import time
	from collections import OrderedDict
	from functools import partial, reduce
	from typing import Any, Dict, List, Tuple

	import pyecharts
	import torch

	from internlm.core.naive_amp import NaiveAMPModel

	mb = 1024 * 1024


	class SimpleMemState:
	"""
	A class to represent the memory state of a model layer.

	Args:
	layer_name (str): The name of the layer.
	layer_mem (int): The memory usage of the layer in bytes.
	"""

	def __init__(self, layer_name: str, layer_mem: int = 0) -> None:
	self.layer_name = layer_name

	# Memory status of the current model layer.
	self._layer_mem: int = layer_mem
	# Total memory status of the model and sub-models, initialized with layer memory.
	self._total_mem: int = self._layer_mem
	# SimpleMemState of sub-models.
	self.sub_model_stats = OrderedDict()

	@property
	def layer_mem(self) -> int:
	"""
	Get the memory usage of the layer.

	Returns:
	int: The memory usage of the layer in bytes.
	"""
	return self._layer_mem

	@layer_mem.setter
	def layer_mem(self, new_layer_mem: int) -> None:
	"""
	Set the memory usage of the layer.

	Args:
	new_layer_mem (int): The new memory usage of the layer in bytes.
	"""
	diff = new_layer_mem - self._layer_mem
	self._layer_mem = new_layer_mem
	self._total_mem += diff

	@property
	def total_mem(self) -> int:
	"""
	Get the total memory usage of the model and sub-models.

	Returns:
	int: The total memory usage in bytes.
	"""
	return self._total_mem

	def add(self, layer_name: str, layer_mem: int = 0, flush: bool = True) -> None:
	"""
	Add a layer to the memory state.

	Args:
	layer_name (str): The name of the layer.
	layer_mem (int, optional): The memory usage of the layer in bytes. Defaults to 0.
	flush (bool, optional): Whether to update the total memory usage. Defaults to True.
	"""
	path = layer_name.split(".")

	target = self.find_layer_state(path, create=True)
	target.layer_mem = layer_mem

	if flush:
	self.update_total_memory()

	def delete(self, layer_name: str, flush: bool = True) -> None:
	"""
	Delete a layer from the memory state.

	Args:
	layer_name (str): The name of the layer.
	flush (bool, optional): Whether to update the total memory usage. Defaults to True.
	"""
	path = layer_name.split(".")
	assert len(path) >= 2, f"Only support deleting non-root layers, layer_name: {layer_name}"

	parent_path = path[0:-1]
	layer = path[-1]
	parent = self.find_layer_state(parent_path)

	if parent is not None and layer in parent.sub_model_stats:
	del parent.sub_model_stats[layer]

	if flush:
	self.update_total_memory()

	def update_total_memory(self) -> None:
	"""
	Update the total memory usage of the model and sub-models.
	"""
	self._total_mem = self._layer_mem

	for stat in self.sub_model_stats.values():
	# Update sub-model status first.
	stat.update_total_memory()
	# Add sub-model total_mem to model total_mem.
	self._total_mem += stat._total_mem

	def find_layer_state(self, path: Tuple[str], create: bool = False) -> "SimpleMemState":
	"""
	Find the memory state of a layer.

	Args:
	path (Tuple[str]): The path to the layer.
	create (bool, optional): Whether to create the layer if it doesn't exist. Defaults to False.

	Returns:
	SimpleMemState: The memory state of the layer.
	"""
	current_node = self

	for _node in path:
	if _node not in current_node.sub_model_stats:
	if not create:
	return None
	# Create a layer node.
	current_node.sub_model_stats[_node] = SimpleMemState(_node)

	current_node = current_node.sub_model_stats[_node]

	return current_node

	def dump(self, prefix: str = "") -> str:
	"""
	Dump the memory state of the model and sub-models.

	Args:
	prefix (str, optional): The prefix to add to the layer names. Defaults to "".

	Returns:
	str: The memory state information.
	"""
	cur_prefix = prefix + "." + self.layer_name if prefix != "" else self.layer_name
	res = f"layer: {cur_prefix}, layer_mem: {self.layer_mem / mb:.2f} MB, total_mem: {self.total_mem / mb:.2f} MB\n"

	for sub_layer in self.sub_model_stats.values():
	res += sub_layer.dump(cur_prefix)

	return res

	def to_json(self, base: int = 1024 * 1024) -> dict:
	"""
	Convert the memory state to a JSON structure.

	Returns:
	dict: The JSON structure of the memory state.
	"""
	children = [child.to_json() for child in self.sub_model_stats.values()]
	if len(children) == 0:
	return {"name": self.layer_name, "value": self.layer_mem // base}
	else:
	return {"name": self.layer_name, "children": children}


	class ActivationMemState:
	"""
	Activation Memory State
	"""

	def __init__(self, num_chunks: int) -> None:
	self._num_chunks = num_chunks

	self.inited: List[bool] = [False for _ in range(num_chunks)]
	self.states: List[SimpleMemState] = [SimpleMemState(f"activations_{idx}") for idx in range(num_chunks)]

	@property
	def total_mem(self) -> int:
	return sum(state.total_mem for state in self.states)

	def dump(self, prefix: str = "") -> str:
	return reduce(lambda x, y: x + y, [state.dump(prefix) for state in self.states])

	def to_json(self, base: int = 1024 * 1024) -> List:
	return [state.to_json(base) for state in self.states]


	def _unpack_naive_wrapper(model: torch.nn.Module) -> Tuple[torch.nn.Module, int]:
	num_chunks = len(model) if isinstance(model, torch.nn.ModuleList) else 1

	if num_chunks > 1:
	model = torch.nn.ModuleList([_model.model if isinstance(_model, NaiveAMPModel) else _model for _model in model])
	else:
	model = model.model if isinstance(model, NaiveAMPModel) else model

	return model, num_chunks


	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:
	"""
	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
	"""
	now = time.time()
	file = f"{self._log_folder}/memory.log"

	if with_options == "all":
	options = ["params", "grads", "os_params", "os_state", "activation_base"]
	else:
	options = with_options.split(",")

	total_mem = (
	self._param_mem_state.total_mem
	+ self._grad_mem_state.total_mem
	+ self._os_params_mem_state.total_mem
	+ self._os_state_mem_state.total_mem
	+ self._activation_mem
	) / mb

	# Generate summary information for memory state
	summary_info = (
	f"total_memory: {total_mem:.2f} MB"
	+ "\n"
	+ f"params_memory: {self._param_mem_state.total_mem / mb:.2f} MB, "
	+ f"grads_memory: {self._grad_mem_state.total_mem / mb:.2f} MB, "
	+ f"os_params_memory: {self._os_params_mem_state.total_mem / mb:.2f} MB, "
	+ f"os_state_memory: {self._os_state_mem_state.total_mem / mb:.2f} MB, "
	+ f"activation_memory: {self._activation_mem / mb:.2f} MB"
	)

	# Generate layout information based on selected options
	layout_info = ""
	if "params" in options:
	layout_info += "params_layout:\n" + self._param_mem_state.dump()
	if "grads" in options:
	layout_info += "grads_layout:\n" + self._grad_mem_state.dump()
	if "os_params" in options:
	layout_info += "os_params_layout:\n" + self._os_params_mem_state.dump()
	if "os_state" in options:
	layout_info += "os_state_layout:\n" + self._os_state_mem_state.dump()
	if "activation_base" in options:
	layout_info += "activation_base_layout:\n" + self._activation_base_mems.dump()

	# Write memory state information to log file
	file_mode = "w" if create else "a"
	with open(file, file_mode, encoding="utf-8") as writer:
	writer.write(
	"Memory State:\n" + f"time: {now - self._record_start_time}\n" + "---summary---\n" + summary_info + "\n"
	)
	if layout_info != "":
	writer.write("---Layout---\n" + layout_info)
	writer.write("\n")

	def step(self) -> None:
	"""
	Update the memory state of the optimizer state.

	Returns:
	None
	"""
	if self._stoped:
	return

	self._remaining_steps -= 1
	if self._remaining_steps == 0:
	self._stoped = True

	# Update os state memory usage
	self._os_state_mem_state = SimpleMemState("os_state_mem")
	self._calc_tensor_group_memory(self._os_state_mem_state, list(self._optimizer.state_dict()["state"].items()))

	if not self._stoped:
	# Do we need to print os_state_layout every time? Is it always constant?
	self.point(with_options="os_state")
	else:
	# Dump memory layout
	self.point(with_options="all")
	# Generate sunburst charts
	self._render_sunburst_chart(self._param_mem_state.to_json()["children"], "params_memory_sunburst")
	self._render_sunburst_chart(self._grad_mem_state.to_json()["children"], "grads_memory_sunburst")
	self._render_sunburst_chart(
	[self._os_params_mem_state.to_json(), self._os_state_mem_state.to_json()],
	"os_memory_sunburst",
	)
	self._render_sunburst_chart(self._activation_base_mems.to_json(), "activation_memory_sunburst")
	# Generate summary sunburst chart
	summary_sunburst_data = [
	{"name": "params", "value": self._param_mem_state.total_mem // mb},
	{"name": "grads", "value": self._grad_mem_state.total_mem // mb},
	{"name": "os_params", "value": self._os_params_mem_state.total_mem // mb},
	{"name": "os_state", "value": self._os_state_mem_state.total_mem // mb},
	{"name": "activation", "value": self._activation_mem_max // mb},
	]

	self._render_sunburst_chart(summary_sunburst_data, "summary_sunburst")

	def _render_sunburst_chart(self, data: Any, name: str) -> None:
	pyecharts.charts.Sunburst(init_opts=pyecharts.options.InitOpts(width="1000px", height="1000px")).add(
	name,
	data_pair=data,
	highlight_policy="ancestor",
	radius=[0, "95%"],
	levels=[
	{},
	{
	"r0": "10%",
	"r": "35%",
	"itemStyle": {"borderWidth": 3},
	"label": {"align": "left"},
	},
	{"r0": "35%", "r": "55%", "label": {"align": "left"}},
	{"r0": "55%", "r": "70%", "label": {"align": "left"}},
	{"r0": "70%", "r": "80%", "label": {"align": "left"}},
	{"r0": "80%", "r": "90%", "label": {"align": "left"}},
	{
	"r0": "90%",
	"r": "92%",
	"label": {"position": "outside", "padding": 3, "silent": False},
	"itemStyle": {"borderWidth": 3},
	},
	],
	).set_global_opts(title_opts=pyecharts.options.TitleOpts(title="CUDA Memory")).set_series_opts(
	label_opts=pyecharts.options.LabelOpts(formatter="{b}")
	).render(
	f"{self._log_folder}/{name}.html"
	)

	def _inner_activation_trace_hook(
	self,
	chunk_id: int,
	layer_name: str,
	model: Any,
	inputs: Any,
	output: torch.Tensor,
	) -> None:
	"""
	Hook function to trace the activation memory usage for a inner layer.

	Args:
	layer_name (str): The name of the layer.
	model (Any): The model.
	inputs (Any): The inputs to the layer.
	output (torch.Tensor): The output tensor.

	Returns:
	None
	"""
	del model, inputs
	assert isinstance(output, torch.Tensor), f"Invalid output type: {type(output)}"

	if self._stoped or self._activation_base_mems.inited[chunk_id]:
	return

	# Delay updating the total_mem of activation_base_mem here, it will be handled in the forward ending hook.
	self._activation_base_mems.states[chunk_id].add(
	layer_name, output.element_size() * output.nelement(), flush=False
	)

	def _activation_trace_hook_forward(self, chunk_id: int, model: Any, inputs: Any, output: torch.Tensor) -> None:
	"""
	Hook function to trace the activation memory usage for a forward pass.

	Args:
	model (Any): The model.
	inputs (Any): The inputs to the model.
	output (torch.Tensor): The output tensor.

	Returns:
	None
	"""
	del model, inputs
	assert isinstance(output, torch.Tensor), f"invalid output type: {type(output)}"

	if self._stoped:
	return

	# Check if the activation memory has been initialized
	if self._activation_base_mems.inited[chunk_id] is False:
	self._activation_base_mems.inited[chunk_id] = True
	# Update the total memory of the activation base memory state
	self._activation_base_mems.states[chunk_id].update_total_memory()
	# Set with_options to "activation_base" to include activation_base_layout in the memory dump
	with_options = "activation_base"
	else:
	with_options = ""

	# Accumulate activation memory usage for each forward pass
	self._activation_mem += self._activation_base_mems.states[chunk_id].total_mem
	if self._activation_mem > self._activation_mem_max:
	self._activation_mem_max = self._activation_mem

	# Trigger a memory record
	self.point(with_options)

	def _activation_tarce_hook_backward(self, chunk_id: int, model: Any, inputs: Any, grad_outputs: Any) -> None:
	"""
	Hook function to trace the activation memory usage for a backward pass.

	Args:
	model (Any): The model.
	inputs (Any): The inputs to the model.
	grad_outputs (Any): The gradients of the outputs.

	Returns:
	None
	"""
	del model, inputs, grad_outputs

	if self._stoped:
	return

	# Release activation memory usage for each backward pass
	self._activation_mem -= self._activation_base_mems.states[chunk_id].total_mem

	# Trigger a memory record
	self.point()

	def _register_activation_trace_hooks(self, chunk_id: int, model_chunk: torch.nn.Module) -> None:
	"""
	Register activation trace hooks for the model and each submodule in the model.
	"""

	# Register inner activation trace hooks for each submodule in the model
	for layer_name, sub_model in model_chunk.named_modules():
	# Register the hook
	if len(sub_model._modules) != 0:
	continue # TODO: in some special cases, we may need some additional configuration to correct

	sub_model.register_forward_hook(partial(self._inner_activation_trace_hook, chunk_id, layer_name))

	# Register a forward hook for the main model to track activation memory usage
	model_chunk.register_forward_hook(partial(self._activation_trace_hook_forward, chunk_id))
	# Register a backward hook for the main model to release activation memory usage
	model_chunk.register_full_backward_hook(partial(self._activation_tarce_hook_backward, chunk_id))

	def _calc_tensor_memory(
	self, root_stat: SimpleMemState, named_tensors: Dict[str, torch.Tensor], require_grad: bool = False
	) -> None:
	"""
	Calculate the memory usage of tensors and update the memory state.

	Args:
	root_stat (SimpleMemState): The root memory state.
	named_tensors (Dict[str, torch.Tensor]): A dictionary containing the named tensors.
	require_grad (bool, optional): Whether to consider tensors with gradients. Defaults to False.

	Returns:
	None
	"""
	for name, tensor in named_tensors:
	if require_grad and not tensor.requires_grad:
	continue

	layer_splits = name.split(sep=".")
	layer_stat = root_stat.find_layer_state(layer_splits, create=True)
	layer_stat.layer_mem = tensor.element_size() * tensor.nelement()

	root_stat.update_total_memory()

	def _calc_tensor_group_memory(self, root_stat: SimpleMemState, tensor_groups: List[Tuple[int, torch.Tensor]]):
	"""
	Calculate the memory usage of a group of tensors.

	Args:
	root_stat (SimpleMemState): The root memory state.
	tensor_groups (List[Tuple[int, torch.Tensor]]): A list of tuples containing the tensor groups.

	Returns:
	None
	"""

	def _normalize_helper(named_tensors: Dict[str, Any]) -> List[Tuple[str, Any]]:
	"""
	Normalize the named tensors.

	Args:
	named_tensors (Dict[str, Any]): The named tensors to normalize.

	Returns:
	List[Tuple[str, Any]]: The normalized named tensors.
	"""
	res = {}

	for name, tensors in named_tensors.items():
	if isinstance(tensors, torch.Tensor):
	res[name] = tensors
	elif isinstance(tensors, (list, tuple)):
	for index, tensor in enumerate(tensors):
	res[f"{name}.{index}"] = tensor
	elif isinstance(tensors, dict):
	for subname, tensor in tensors.items():
	res[f"{name}.{subname}"] = tensor
	else:
	raise TypeError(f"unsupported normalize value type: {type(tensors)}")

	return list(res.items())

	def _value_check(tensor_or_tensors):
	"""
	Check if the input is a tensor or a collection of tensors.

	Args:
	tensor_or_tensors (Any): The input to check.

	Returns:
	bool: True if the input is a tensor or a collection of tensors, False otherwise.
	"""
	if torch.is_tensor(tensor_or_tensors):
	return True
	elif isinstance(tensor_or_tensors, (list, tuple)) and all(torch.is_tensor(x) for x in tensor_or_tensors):
	return True
	elif isinstance(tensor_or_tensors, dict) and all(torch.is_tensor(x) for x in tensor_or_tensors.values()):
	return True
	else:
	return False

	# Calculate the memory usage of a group of tensors.
	for idx, tensors in tensor_groups:
	# Normalize the named tensors
	named_tensors = {f"{idx}.{k}": v for k, v in tensors.items() if _value_check(v)}
	named_tensors = _normalize_helper(named_tensors)
	# Calculate the memory usage of the tensors and update the memory state
	self._calc_tensor_memory(root_stat, named_tensors)


	if __name__ == "__main__":

	class SimpleModel(torch.nn.Module):
	"""
	A simple model with three linear layers.

	Args:
	skip_layer2 (bool, optional): Whether to skip layer2. Defaults to False.
	"""

	def __init__(self, skip_layer2: bool = False):
	super().__init__()
	self.layer1 = torch.nn.Linear(5120, 5120, True)
	self.layer3 = torch.nn.Linear(5120, 5120, False)

	if skip_layer2:
	self.layer2 = None
	else:
	self.layer2 = SimpleModel(skip_layer2=True)

	def forward(self, inputs: torch.Tensor) -> torch.Tensor:
	"""
	Forward pass of the model.

	Args:
	inputs (torch.Tensor): The input tensor.

	Returns:
	torch.Tensor: The output tensor.
	"""
	output1 = self.layer1(inputs)
	if self.layer2 is not None:
	output2 = self.layer2(output1)
	else:
	output2 = output1
	output = self.layer3(output2)

	return output

	def _simple_schedule(_num_chunks, _model_chunks, _input) -> torch.Tensor:
	if _num_chunks > 1:
	_output = _input
	for _model_chunk in _model_chunks:
	_output = _model_chunk(_output)
	else:
	_output = _model_chunks(_input)

	return _output

	# num_chunks config
	_num_chunks = 1

	# init model and optimizer
	if _num_chunks > 1:
	_chunks = [SimpleModel(skip_layer2=idx % 2 == 0) for idx in range(_num_chunks)]
	_model = torch.nn.ModuleList(_chunks).cuda()
	else:
	_model: torch.nn.Module = SimpleModel().cuda()
	_optimizer = torch.optim.Adam(_model.parameters())

	# init profiler
	profiler = SimpleMemoryProfiler(_model, _optimizer, "./test_simple_memory_profiler", total_steps=1)

	_optimizer.zero_grad()

	# inputs
	x1 = torch.randn((128, 5120)).cuda()
	x2 = torch.randn((128, 5120)).cuda()
	# forward
	out1 = _simple_schedule(_num_chunks, _model, x1)
	out2 = _simple_schedule(_num_chunks, _model, x2)
	# backward
	out1.mean().backward()
	out2.mean().backward()

	_optimizer.step()

	# Update the optimizer state memory usage and record the memory state
	profiler.step()