Diff Coverage

# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Walkthrough of the EvaluatorV2 memory estimation API, see README.md."""
from typing import Any

from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2 import EvaluatorV2
from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.hooks.template import Template
from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger


def my_attn_num_param(ccfg: Any, ctx: Any) -> float:
    """Attention parameter count that overrides the default formula."""
    del ctx
    return 10 * ccfg.h * ccfg.h


def custom(ccfg: Any) -> None:
    """Cost model variables that override the parsed configuration."""
    ccfg.bytes_compute = 1
    ccfg.s = 1024
    ccfg.n_attMM = 5

    ccfg.s = 1024
    ccfg.n_attMM = 5


def main() -> None:
    """Estimate, inspect and customize the memory of the bundled test cases."""
    # Instantiate evaluator with a model configuration,
    #  log_level=0 removes warning messages
    e = EvaluatorV2("./test_cases/mixtral/default.yaml", log_level=0)

    # Check all defined node type
    logger.output("%s", list(LayerType))

    # Estimate peak memory (in Megabytes)
    peak_mem = e.estimate_peak(verbose=True)
    # Check whether estimation fits in device's max memory
    e.mem_fit(peak_mem)

    # Estimate static memory of a specific pipeline stage (in Megabytes)
    logger.output("%s", e.static_mem_stage(1))
    # Estimate dynamic memory of a specific pipeline stage (in Megabytes)
    logger.output("%s", e.dynamic_mem_stage(1))
    # Estimate static memory of a specific layer and stage (in Megabytes)
    logger.output("%s", e.static_mem_layer(LayerType.FULL_REC_LAYER, 1))
    # Estimate dynamic memory of a specific layer and stage (in Megabytes)
    logger.output("%s", e.dynamic_mem_layer(LayerType.FULL_REC_LAYER, 1))
    # Retrieve the memory estimation logs of a specific stage (in Megabytes)
    logger.output("%s", e.logs_mem_stage(1))
    # Fetch memory insights from each pipeline stage
    logger.output("%s", e.estimate_peak_insight())
    # PPB Input
    logger.output("%s", e.estimate_layer_memory())

    # Inspect a specific stage (here is the first one)
    e.estimate_peak(spec_stage_id=0, verbose=True)

    # Plot
    e.estimate_peak(plot=True)

    e = EvaluatorV2("./test_cases/deepseek3/default.yaml", log_level=0)

    # Overwriting context function
    e.set_attn_eval_fun(num_p=my_attn_num_param)

    # Overwriting a training feature
    e.set_passes(swap_os=True)

    # Overwriting cost model variables
    e.set_ccfg(custom)

    # Overwriting strategy
    logger.output("%s", e.get_strategy())
    e.set_strategy(dp=8, tp=8, m=128)
    logger.output("%s", e.get_strategy())

    e.estimate_peak(verbose=True)
    # Inspect ccfg object (cost model variables)
    e.print_ccfg()
    # Inspect ctx object (evaluation variables and functions)
    e.print_ctx()

    # Load a hook class
    # ... when declaring an Evaluator
    e = EvaluatorV2(
        "./test_cases/deepseek3/default.yaml", log_level=0, hook_cls=Template()
    )
    # ... by using load_hook_cls()
    e.load_hook_cls(Template())


if __name__ == "__main__":
    main()

                    solver, group_name, inter, self._get_body_layer_by_name(group_name).recompute_considered_,
                )
                for stage_id, rec_counts in enumerate(chunk_stage_rec):
                    for rec_type, count in rec_counts:
                        stage_partition[inter * pp + stage_id].extend(
                            (layer_id, rec_type) for layer_id in range(current_layer_id, current_layer_id + count)
                        )
                        current_layer_id += count

        if step_type == MetaStepType.DATA_LOAD:
            if stage_index <= 0:
                self._exec_data_load(micro_index, arg_mbs, kwarg_mbs)
        elif step_type == MetaStepType.DATA_SEND:
            self._exec_data_send(stage_index, micro_index, arg_mbs, kwarg_mbs)
        elif step_type == MetaStepType.DATA_RECV:
            self._exec_data_recv(stage_index, micro_index, arg_mbs, kwarg_mbs)

    def _exec_data_load(self, micro_index: int, arg_mbs: list, kwarg_mbs: list) -> None:
        """Load the next micro-batch onto stage 0's device and stage it as ``micro_index``'s input."""
        device = self.stages[0].device

        """Load the next micro-batch onto stage 0's device and stage it as ``micro_index``'s input."""
        device = self.stages[0].device
        micro_batch = next(self.data_iterator)
        if isinstance(micro_batch, list):
            micro_batch = micro_batch[micro_index]
        micro_batch = {
            key: (value.to(device, non_blocking=True) if hasattr(value, "to") else value)
            for key, value in micro_batch.items()
        }

        self.last_local_tokens += n_valid

        micro_batch["targets"] = targets
        for key in self.kwargs_batch_dim:
            if key in kwarg_mbs[micro_index].keys():
                kwarg_mbs[micro_index][key] = torch.cat(
                    [kwarg_mbs[micro_index][key], micro_batch[key]], dim=0)
            else:
                kwarg_mbs[micro_index][key] = micro_batch[key]
        arg_mbs[micro_index] = [input_ids]

    def _exec_data_send(self, stage_index: int, micro_index: int, arg_mbs: list, kwarg_mbs: list) -> None:
        """Send ``micro_index``'s data keys to the stage that consumes them."""
        if getattr(self, "_data_dst", False):
            dst_stage_idx = self._data_dst[stage_index][micro_index]
        else:
            dst_stage_idx = self.stages[0].dst_stage
        dst = self.stages[0]._global_rank(dst_stage_idx)  # pylint: disable=protected-access
        # One meta round-trip per DATA_SEND instead of K; the matching
        # DATA_RECV unpacks the list in the same order.
        metas = []
        tensors = []
        for key in self._data_keys:
            tensor = (arg_mbs[micro_index][0] if key == "input_ids"
                      else kwarg_mbs[micro_index][key])
            metas.append([tuple(tensor.shape), tensor.dtype])
            tensors.append(tensor)
        dist.send_object_list(metas, dst)
        handles = [dist.isend(t, dst) for t in tensors]
        self._wait_p2p(handles)

    def _exec_data_recv(self, stage_index: int, micro_index: int, arg_mbs: list, kwarg_mbs: list) -> None:
        """Receive ``micro_index``'s data keys and stage them as its input."""
        if getattr(self, "_data_src", False):
            src_stage_idx = self._data_src[stage_index][micro_index]
        else:
            src_stage_idx = self.stages[0].src_stage
        src = self.stages[0]._global_rank(src_stage_idx)  # pylint: disable=protected-access
        device = self.stages[0].device
        # One recv_object_list unpacks every key's (shape, dtype) at
        # once, matching the packed DATA_SEND above.
        metas: list = [None] * len(self._data_keys)
        dist.recv_object_list(metas, src)
        handles = []
        for key, meta in zip(self._data_keys, metas):
            shape, dtype = meta
            buffer = torch.empty(shape, dtype=dtype, device=device)
            handles.append(dist.irecv(buffer, src))
            if key == "input_ids":
                arg_mbs[micro_index] = [buffer]
            else:
                kwarg_mbs[micro_index][key] = buffer
        self._wait_p2p(handles)

    def _exec_pipeline_swap_step(self, cur_step, arg_mbs, kwarg_mbs):
        """Execute a pipeline activation-swap control step."""
        if self._swap_session is None: