Diff Coverage

recomputation policies across pipeline parallel stages.
"""
# pylint: disable=invalid-name,undefined-all-variable

from importlib import import_module as _import_module
import sys as _sys

_sys.modules.setdefault("sapp_ppb", _sys.modules[__name__])

__all__ = [
    "SappPipeline",
    "choose_interleave",
    "flatten",
    "SappSolver",

    "run",
    "main",
]

_EXPORTS = {
    "SappPipeline": ".sapp.sapp_pipeline",
    "choose_interleave": ".sapp.sapp_pipeline",
    "flatten": ".sapp.sapp_pipeline",
    "SappSolver": ".sapp.sapp_solver",

    "main": ".run_pipeline_balance",
}


def __getattr__(name):
    """Lazily import public SAPP-PPB interfaces."""
    if name not in _EXPORTS:
        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

    module = _import_module(_EXPORTS[name], __name__)
    value = getattr(module, name)
    globals()[name] = value
    return value

# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""CLI entrypoint for SAPP-PPB pipeline balancing."""
import argparse
import os
import sys

from sapp_ppb.sapp.sapp_pipeline import SappPipeline
from sapp_ppb.utils import interactive
from sapp_ppb.utils.compute_memory import compute_memories
from sapp_ppb.utils.config import initialize_layer_json
from sapp_ppb.utils.layer import generate_layers_list
from sapp_ppb.utils.logger import logger


def _str2bool(value: str) -> bool:
    """Parse a truthy string value coming from ``argparse``."""
    return str(value).lower() in ('true', '1', 'yes')


def build_arg_parser() -> argparse.ArgumentParser:
    """Build the argument parser for the pipeline-balance CLI.

    Returns:
        Configured :class:`argparse.ArgumentParser` instance.

    Returns:
        Configured :class:`argparse.ArgumentParser` instance.
    """
    parser = argparse.ArgumentParser(
        prog='SAPP AutoBalancing',
        description='Balance layers onto pipeline stages, considering recomputation and interleaving',
        epilog='')

        description='Balance layers onto pipeline stages, considering recomputation and interleaving',
        epilog='')

    # Pipeline info
    parser.add_argument('-s', '--stage', type=int, default=4, help="Number of stages")
    parser.add_argument('-mb', '--micro_batch', type=int, default=4, help="Number of micro batch")
    parser.add_argument('-i', '--interleave_degree', type=int, default=1, help="Interleave level")

    # Memory size
    parser.add_argument('-mem', '--max_memory', type=int, default=56000,
                        help="Maximum memory available (MB)")
    parser.add_argument('-lm', '--less_memory', type=_str2bool, default=False,
                        help="Compute Memory with 'Less Memory interleave' option")
    parser.add_argument('-dual', '--dualpipe_v', type=_str2bool, default=False,
                        help="Compute Memory with 'DualpipeV' option")
    parser.add_argument('-mc', '--constant_memory', type=int, default=0,
                        help="Constant memory per stages")

    parser.add_argument('-o', '--output_folder', type=str, default="./output",
                        help="output files location")

    # Model info
    parser.add_argument('-m', '--model_name', type=str, default="model_name", help="")

    # Search time
    parser.add_argument('-t', '--time_limit', type=int, default=90,
                        help="Limitation on searching time")

    # Optimization level
    parser.add_argument('-O', '--optimization_level', type=int, default=1,
                        help="Defines optimization level when Stage (S) = Micro Batch number (M). "
                             "0 for same approach as M > S. "
                             "1 (default) generally better. "
                             "2 better for memory constrained cases.")

                             "1 (default) generally better. "
                             "2 better for memory constrained cases.")

    # Simulate naive or manual config
    parser.add_argument('-naive', '--simulate_naive', type=_str2bool, default=False,
                        help="Simulate naive configs")
    parser.add_argument('-manual', '--manual_config', type=str, default=None,
                        help="Path of manual config")

    # Layer info
    parser.add_argument('-lf', '--layer_folder', type=str, default="./layers/",
                        help="Path to the layer folder")
    parser.add_argument('-dump', '--dump_layer', type=_str2bool, default=False,
                        help="Dump the layers")

    # For Computation of memory
    parser.add_argument('-mf', '--memory_folder', type=str, default="./memory/",
                        help="Path to the profiler memory folder")

    # For Initialization
    parser.add_argument('-init', '--init', type=str, default=None,
                        help="Path to the init file")

    # Computation argument
    parser.add_argument('-cm', '--compute_memory', type=_str2bool, default=False,
                        help="Parse Mindspore log to generate MEMORY of the layer (unavailable)")
    parser.add_argument('-exec', '--exec', type=_str2bool, default=True,
                        help="Compute solver")
    return parser


def _resolve_path(base_dir: str, path: str) -> str:
    """Return ``path`` resolved relative to ``base_dir`` unless it is already absolute."""
    if os.path.isabs(path):
        return path
    return os.path.join(base_dir, path)


def run(args: argparse.Namespace, base_dir: str) -> None:
    """Execute the pipeline balancing workflow for the given arguments.

    Args:
        args (argparse.Namespace): Parsed CLI arguments.

    Args:
        args (argparse.Namespace): Parsed CLI arguments.
        base_dir (str): Directory used to resolve relative input / output paths.
    """
    if args.init:
        init_file = _resolve_path(base_dir, args.init)
        initialize_layer_json(args.model_name, init_file)

    output_folder = _resolve_path(base_dir, args.output_folder)
    os.makedirs(output_folder, exist_ok=True)

    manual_config = None
    if args.manual_config:
        candidate = _resolve_path(base_dir, args.manual_config)
        if candidate.endswith(('yaml', 'yml')):
            manual_config = candidate

    layers = generate_layers_list(args.layer_folder, args.model_name)
    if args.compute_memory:
        layers = compute_memories(layers=layers, memory_folder=args.memory_folder,
                                  number_of_stage=args.stage)
    for layer in layers:
        logger.output("%s", layer)

    if args.dump_layer:
        for layer in layers:
            layer.dump()

    pipe = SappPipeline(model_name=args.model_name, num_of_stage=args.stage,
                        num_of_micro_batch=args.micro_batch, max_memory=args.max_memory,
                        layers=layers, num_of_interleave=args.interleave_degree,
                        vpp_less_memory=args.less_memory, dual=args.dualpipe_v,
                        constant_memory=args.constant_memory,

                        vpp_less_memory=args.less_memory, dual=args.dualpipe_v,
                        constant_memory=args.constant_memory,
                        optimization_level=args.optimization_level)

    pipe.construct_problem(solver="pulp")

    if args.exec:
        pipe.solve_problem(time_limit=args.time_limit, dump_folder=output_folder)
        pipe.print_yaml_results()
        total_time = pipe.simulate(show=True, file_name=os.path.join(output_folder, "result.svg"))

        logger.output("total_time: %d", total_time)
        logger.output("time: %s", pipe.get_time())
        logger.output("mem_par: %s", pipe.get_memory_parameter())
        logger.output("mem_act: %s", pipe.get_memory_activation())

        if manual_config:
            logger.output("Simulating manual configs")
            pipe.simulate_comparison(manual_config, output_folder)
        if args.simulate_naive:
            logger.output("Simulating naive configs")
            pipe.simulate_naive(layers, output_folder)
    elif manual_config:
        logger.output("Simulating manual configs")
        pipe.simulate_only_manual(manual_config, output_folder)


def main() -> None:
    """Entry point invoked when the module is run as a script."""
    if len(sys.argv) == 1:
        interactive.main()
        return
    parser = build_arg_parser()
    args = parser.parse_args()
    run(args, base_dir=os.path.dirname(os.path.abspath(__file__)))


if __name__ == "__main__":
    main()

# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""High-level orchestrator around :class:`SappSolver`: build, solve, simulate, export YAML."""
import os
import sys
from typing import Any, Dict, List, Optional, Union

import matplotlib.pyplot as plt
import yaml

import sapp_ppb.simulator.pp_simulator as sim
import sapp_ppb.utils.recompute as Recompute
from sapp_ppb.sapp.sapp_solver import SappSolver
from sapp_ppb.utils.check_rules import check_yaml_depth_before_loading
from sapp_ppb.utils.layer import Layer, filter_layer_type
from sapp_ppb.utils.logger import logger


class SappPipeline:
    """pipeline balancer"""

    def __init__(
            self,
            model_name: str,
            num_of_stage: int,
            num_of_micro_batch: int,

                seqpp. Default: ``None``.
            seq_split_num (int, optional): Number of sequence splits; ``>1`` enables sequence pipeline.
                Default: ``1``.
        """
        self.model_name_ = model_name
        self.num_of_stage_ = num_of_stage
        self.num_of_micro_batch_ = num_of_micro_batch
        self.num_of_interleave_ = num_of_interleave
        self.max_memory_ = max_memory
        self.vpp_less_memory_ = vpp_less_memory
        # Add arg dual_
        self.dual_ = dual
        self.constant_memory_ = constant_memory
        self.optimization_level = optimization_level
        self.extracted_training_params_ = extracted_training_params
        self.seq_split_num_ = seq_split_num
        self.seqpipe_ = self.seq_split_num_ > 1
        # logger.output("seq chunk: %s",self.seq_split_num_)

        self.problem_ = None
        self.layers_ = layers
        self.layers_sorted_ = {
            Layer.type_enum.HEAD: filter_layer_type(layers,
                                                    Layer.type_enum.HEAD),
            Layer.type_enum.BODY: filter_layer_type(layers,
                                                    Layer.type_enum.BODY),

            Layer.type_enum.TAIL: filter_layer_type(layers,
                                                    Layer.type_enum.TAIL),
        }

    def has_some_memory_info(self) -> bool:
        """Check if there is all information for memory constraint."""
        return self.problem_.has_some_memory_info()

    def construct_problem(self, solver: str = "pulp") -> None:
        """Construct the underlying ILP problem using the requested solver backend."""
        if solver == "pulp":
            self.problem_ = self._construct_problem_pulp_()
        elif solver == "other":
            logger.warning(
                "No other solver available..., automatically switch to pulp!!!"
            )
            self.problem_ = self._construct_problem_pulp_()
        else:
            logger.warning(
                "No other solver available..., automatically switch to pulp!!!"
            )
            self.problem_ = self._construct_problem_pulp_()

    def solve_problem(self, time_limit: int = 90, dump_folder: Optional[str] = None) -> None:
        """Solve the ILP, optionally dumping the LP model into ``dump_folder``."""
        self.problem_.solve(time_limit, dump_folder)

    def get_result(self) -> dict[str, list[list[str]]]:
        """Get result distribution of the solution (compact form)."""
        return self.problem_.result()

    def get_memory_activation(self) -> list[float]:
        """Get the activation memory per stage for simulator."""
        return self.problem_.get_simulator_memory_activation()

    def get_memory_parameter(self) -> list[float]:
        """Get the parameter memory per stage for simulator."""
        return self.problem_.get_simulator_memory_parameter()

    def get_fw_time(self) -> list[float]:
        """Get the forward time per stage for simulator."""
        time = self.problem_.get_simulator_forward_time()
        return time

    def get_recompute_time(self) -> list[float]:
        """Get the recompute time per stage for simulator."""
        time = self.problem_.get_simulator_recompute_time()
        return time

    def get_time(self) -> list[float]:
        """Get the time per stage for simulator."""
        return self.problem_.get_simulator_time()

    def naive_layer_per_stage(self,
                              layer_num: int,
                              num_of_interleave: int = 1) -> List[List[int]]:
        """Return the naive layer-to-stage assignment (``layer_num`` evenly split)."""
        logger.output("layer_num = %s", layer_num)
        layer_count = layer_num // (self.num_of_stage_ * num_of_interleave)
        return [[layer_count] * self.num_of_stage_ for _ in range(num_of_interleave)]

    def print_yaml_results(self) -> None:
        """Log the solver output in the MindFormers YAML schema."""

        for layer in self.layers_sorted_[Layer.type_enum.BODY]:
            nass = self.naive_layer_per_stage(layer.nb_layer_,
                                              self.num_of_interleave_)
            yaml_format = Recompute.yaml_from_internal(
                self.num_of_interleave_,
                self.num_of_stage_,
                self.problem_.variables_[layer.name_],
                nass,

                self.num_of_stage_,
                self.problem_.variables_[layer.name_],
                nass,
            )
            logger.output("layer-to-stage assignment baseline is \n\t%s", nass)
            yaml_results = "\nTo put in yaml configuration:"
            for y, v in yaml_format.items():
                yaml_results += f"\n\t{y}: {v}"
            logger.output(yaml_results)

    def get_manual_memory_activation(
            self,
            each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]],
            interleave_num: int = 1) -> List[List[float]]:
        """Return the per-stage activation memory for a user-supplied layer assignment."""
        memory_active = []
        if self.has_some_memory_info():
            for inter in range(interleave_num):
                memory_active.append([])
                for stage in range(self.num_of_stage_):
                    memory_active[inter].append(sum(
                        each_layer_per_recompute[layer][rec][inter][stage] *
                        layer.memory_activation_rec_[rec]
                        for layer in self.layers_sorted_[Layer.type_enum.BODY]
                        for rec in Recompute.TYPE

                        for layer in self.layers_sorted_[Layer.type_enum.BODY]
                        for rec in Recompute.TYPE
                        if rec not in Recompute.get_unused_list(each_layer_per_recompute[layer])
                        and each_layer_per_recompute[layer][rec][inter][stage] > 0))
        return memory_active

    def get_manual_memory_parameter(
            self,
            each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]],
            interleave_num: int = 1) -> List[List[float]]:
        """Return the per-stage parameter memory for a user-supplied layer assignment."""
        memory_param_stage = [0] * self.num_of_stage_
        for inter in range(interleave_num):
            for stage in range(self.num_of_stage_):
                memory_param_stage[stage] += sum(
                    each_layer_per_recompute[layer][rec][inter][stage] *
                    layer.memory_parameter_ for rec in Recompute.TYPE
                    for layer in self.layers_sorted_[Layer.type_enum.BODY]
                    if layer.memory_parameter_ is not None

                    for layer in self.layers_sorted_[Layer.type_enum.BODY]
                    if layer.memory_parameter_ is not None
                    and rec not in Recompute.get_unused_list(each_layer_per_recompute[layer])
                    and each_layer_per_recompute[layer][rec][inter][stage] > 0)
        for head in self.layers_sorted_[Layer.type_enum.HEAD]:
            if head.memory_parameter_ is not None:
                memory_param_stage[0] += head.memory_parameter_
        for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
            if tail.memory_parameter_ is not None:
                memory_param_stage[self.num_of_stage_ -
                                   1] += tail.memory_parameter_
        memory_param = [memory_param_stage] * interleave_num
        return memory_param

    def get_manual_time(
            self,
            each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]],
            interleave_num: int = 1) -> List[List[float]]:
        """Return the per-stage execution time for a user-supplied layer assignment."""
        time = []
        for i in range(interleave_num):
            time.append([])
            for s in range(self.num_of_stage_):
                time[i].append(0)
                for layer in self.layers_sorted_[Layer.type_enum.BODY]:
                    for r in Recompute.TYPE:
                        if each_layer_per_recompute[layer][r][i][s] > 0:
                            time[i][s] += each_layer_per_recompute[layer][r][i][s] * (
                                layer.forward_time_ +
                                layer.backward_time_rec_[r])

        for head in self.layers_sorted_[Layer.type_enum.HEAD]:
            time[0][0] += head.time_
        for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
            time[interleave_num - 1][self.num_of_stage_ - 1] += tail.time_
        return time

    def get_manual_fw_time(
            self,
            each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]],
            interleave_num: int = 1) -> List[List[float]]:
        """Return the per-stage forward time for a user-supplied layer assignment."""
        time = []
        for i in range(interleave_num):
            time.append([])
            for s in range(self.num_of_stage_):
                time[i].append(0)
                for layer in self.layers_sorted_[Layer.type_enum.BODY]:
                    for r in Recompute.TYPE:
                        if (r not in Recompute.get_unused_list(each_layer_per_recompute[layer])
                            and each_layer_per_recompute[layer][r][i][s] > 0):
                            time[i][s] += each_layer_per_recompute[layer][r][i][s] * (
                                layer.forward_time_)
        for head in self.layers_sorted_[Layer.type_enum.HEAD]:
            time[0][0] += head.time_
        for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
            time[interleave_num - 1][self.num_of_stage_ - 1] += tail.time_
        return time

    def get_manual_recompute_time(
            self,
            each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]],
            interleave_num: int = 1) -> List[List[float]]:
        """Return the per-stage recompute-only time for a user-supplied layer assignment."""
        logger.output("each_layer_per_recompute = %s", each_layer_per_recompute)
        time_all_rec = []
        time_no_rec = []
        for i in range(interleave_num):
            time_all_rec.append([])
            time_no_rec.append([])
            for s in range(self.num_of_stage_):
                time_all_rec[i].append(0)
                time_no_rec[i].append(0)
                for layer in self.layers_sorted_[Layer.type_enum.BODY]:
                    self._add_manual_recompute_time(
                        each_layer_per_recompute, layer, i, s, time_all_rec, time_no_rec)

        return [[r - n for r, n in zip(ar, nr)]
                for ar, nr in zip(time_all_rec, time_no_rec)]

    def _add_manual_recompute_time(self, each_layer_per_recompute, layer, interleave, stage,
                                   time_all_rec, time_no_rec):
        """Accumulate recompute time for a single layer and stage."""
        logger.output("backward_time_rec_(%s) = %s", layer, layer.backward_time_rec_)
        unused_rec = Recompute.get_unused_list(each_layer_per_recompute[layer])
        for rec in Recompute.TYPE:
            layer_num = each_layer_per_recompute[layer][rec][interleave][stage]
            if rec in unused_rec or layer_num <= 0:
                continue
            if layer.backward_time_rec_[rec] is None:
                raise ValueError("No backward tme is specified for this "
                                 "recomputation. Recomputation "
                                 f"'{Recompute.YAML_NAME[rec]}' is likely not considered")
            logger.output("r = %s; i = %s; s = %s", rec, interleave, stage)
            time_all_rec[interleave][stage] += layer_num * layer.backward_time_rec_[rec]
            time_no_rec[interleave][stage] += layer_num * layer.backward_time_rec_[Recompute.TYPE.NONE]

    def simulate(self, show: bool = True, file_name: Optional[str] = None,
                 sub_fig: Optional[plt.Figure] = None) -> float:
        """Run the simulator on the solved schedule and return its estimated total time."""
        forward_time = self.get_fw_time()
        recompute_overhead = self.get_recompute_time()
        stage_mem_par = 0
        stage_mem_act = 0
        if self.has_some_memory_info():
            stage_mem_par = self.get_memory_parameter()
            stage_mem_act = self.get_memory_activation()

        return self.simulation(
            forward_time,
            recompute_overhead,
            stage_mem_par,
            stage_mem_act,

            file_name,
            sub_fig
        )

    def simulate_naive(self, layers: List[Layer], output_folder: str) -> None:
        """Simulate the naive (even) layer-to-stage assignments for sanity comparison."""
        num_layers = 0
        rec_considered = {}
        for layer in layers:
            if layer.type_ == Layer.type_enum.BODY:
                num_layers = layer.nb_layer_
                rec_considered = layer.recompute_considered_

        all_recomp = {"offset": 0}
        no_recomp = {"offset": 0}
        for rec in [Recompute.TYPE.FULL, Recompute.TYPE.SLCT, Recompute.TYPE.COMM]:
            if rec_considered.get(rec, False):
                all_recomp[Recompute.YAML_NAME[rec]] = True
                no_recomp[Recompute.YAML_NAME[rec]] = False

        self.simulate_yaml(
            yaml_format=all_recomp,
            show=True,
            interleave_num=self.num_of_interleave_,
            file_name=os.path.join(output_folder,

            file_name=os.path.join(output_folder,
                                   "result_naive_all_recomp.svg"),
        )

        if num_layers % self.num_of_stage_ == 0:
            self.simulate_yaml(
                yaml_format=no_recomp,
                show=True,
                interleave_num=self.num_of_interleave_,
                file_name=os.path.join(output_folder,

                file_name=os.path.join(output_folder,
                                       "result_naive_no_recomp.svg"),
            )
        else:
            logger.warning("num layer cannot be divided by num stage")

    def simulate_comparison(self, manual_config_file: str, output_folder: str) -> None:
        """Render side-by-side automatic vs manual simulations for every entry in the YAML."""
        with open(manual_config_file, encoding="utf-8") as fp:
            check_yaml_depth_before_loading(fp)
            fp.seek(0)
            data = yaml.safe_load(fp)
        yaml_data = {}
        for manual in data.values():
            yaml_data[Recompute.OFFSET] = manual.get(Recompute.OFFSET)
            if isinstance(yaml_data[Recompute.OFFSET], list) and all(
                    isinstance(item, int) for item in yaml_data[Recompute.OFFSET]):
                yaml_data[Recompute.OFFSET] = [yaml_data[Recompute.OFFSET]]

            for rec in Recompute.YAML_NAME.values():
                yaml_data[rec] = manual.get(rec)
                if isinstance(yaml_data[rec], list) and all(
                        isinstance(item, int) for item in yaml_data[rec]):
                    yaml_data[rec] = [yaml_data[rec]]
            interleave_num = manual.get("interleave_num",
                                        self.num_of_interleave_)
            show = manual.get("show", False)
            file_name = manual.get("file_name")
            full_file_name = os.path.join(output_folder,
                                          file_name) if (file_name) else None

            fig = plt.figure(figsize=(24, 8))
            sub_figs = fig.subfigures(1, 2, wspace=0.07)
            sub_figs[0].suptitle('Automatic', fontsize='x-large')
            self.simulate(show=False, file_name=os.path.join(output_folder, "Auto_" + file_name), sub_fig=sub_figs[0])

            sub_figs[1].suptitle('Manual', fontsize='x-large')
            self.simulate_yaml(yaml_data, False, interleave_num, full_file_name, sub_figs[1])
            plt.savefig(os.path.join(output_folder, "Comparison_" + file_name))
            if show:
                plt.show()

    def simulate_only_manual(self, manual_config_file: str, output_folder: str) -> None:
        """Render only the manual simulation for every entry in ``manual_config_file``."""
        with open(manual_config_file, encoding="utf-8") as fp:
            check_yaml_depth_before_loading(fp)
            fp.seek(0)
            data = yaml.safe_load(fp)
        yaml_data = {}
        for manual in data.values():
            yaml_data[Recompute.OFFSET] = manual.get(Recompute.OFFSET)
            if isinstance(yaml_data[Recompute.OFFSET], list) and all(
                    isinstance(item, int) for item in yaml_data[Recompute.OFFSET]):
                yaml_data[Recompute.OFFSET] = [yaml_data[Recompute.OFFSET]]

            for rec in Recompute.YAML_NAME.values():
                yaml_data[rec] = manual.get(rec)
                if isinstance(yaml_data[rec], list) and all(
                        isinstance(item, int) for item in yaml_data[rec]):
                    yaml_data[rec] = [yaml_data[rec]]
            interleave_num = manual.get("interleave_num",
                                        self.num_of_interleave_)
            show = manual.get("show", False)
            file_name = manual.get("file_name")
            full_file_name = os.path.join(output_folder,
                                          file_name) if (file_name) else None

            fig = plt.figure(figsize=(12, 8))
            self.simulate_yaml(yaml_data, False, interleave_num, full_file_name, fig)
            plt.savefig(os.path.join(output_folder, "manual_file_" + file_name))
            if show:
                plt.show()

    def simulate_yaml(self, yaml_format: Dict[str, Any], show: bool = True,
                      interleave_num: int = 1,
                      file_name: Optional[str] = None,
                      sub_fig: Optional[plt.Figure] = None) -> float:
        """Simulate a manual pipeline configuration encoded as a YAML-compatible dict."""
        layer_num = 0
        for layer in self.layers_sorted_[Layer.type_enum.BODY]:
            layer_num += layer.nb_layer_
        nass = self.naive_layer_per_stage(layer_num,
                                          num_of_interleave=interleave_num)
        layer_per_recompute = Recompute.internal_from_yaml(
            interleave_num, self.num_of_stage_, yaml_format, nass)
        each_layer_per_recompute = self.split_layer_per_recompute(layer_per_recompute)
        return self.simulate_manual(
            each_layer_per_recompute,
            show,
            interleave_num=interleave_num,
            file_name=file_name,

    ##                                                                   ##
    ##                      Print Solver Model                           ##
    ##                                                                   ##
    #######################################################################
    def _calculate_activation_memory(self, each_layer_per_recompute, v, s):
        """Calculate activation memory for next and current stage"""
        act_mem_next = 0
        act_mem_curr = 0

        for layer in self.layers_sorted_[Layer.type_enum.BODY]:
            for rec in Recompute.TYPE:
                if self.problem_.recompute_considered_[rec]:
                    if each_layer_per_recompute[layer][rec][v + 1][s] > 0:  # next
                        act_mem_next += (each_layer_per_recompute[layer][rec][v + 1][s] *
                                         layer.memory_activation_rec_[rec])
                    if each_layer_per_recompute[layer][rec][v][s] > 0:    # current
                        act_mem_curr += (each_layer_per_recompute[layer][rec][v][s] *
                                         layer.memory_activation_rec_[rec])

        return act_mem_next, act_mem_curr

    def _compute_parameter_memory_manually_solver(self, each_layer_per_recompute, s, interleave_num=1):
        """Solver memory model: parameter memory"""
        param_mem = 0
        for layer in self.layers_sorted_[Layer.type_enum.BODY]:
            if layer.memory_parameter_ is not None:
                param_mem += self._calculate_layer_parameter_memory(
                    layer, each_layer_per_recompute[layer], s, interleave_num)
        return param_mem

    def _calculate_layer_parameter_memory(self, layer, layer_per_recompute, s, interleave_num):
        """Calculate parameter memory for a single layer"""
        layer_mem = 0
        for inter in range(interleave_num):
            for rec in Recompute.TYPE:
                if self.problem_.recompute_considered_[rec]:
                    if layer_per_recompute[rec][inter][s] > 0:
                        layer_mem += layer_per_recompute[rec][inter][s] * layer.memory_parameter_
        return layer_mem

    def _calculate_activation_memory_solver(self, each_layer_per_recompute, s, interleave_num, activation_nums):
        """Calculate activation memory for a given stage"""
        act_mem = 0
        for layer in self.layers_sorted_[Layer.type_enum.BODY]:
            for inter in range(interleave_num):
                for rec in Recompute.TYPE:
                    if self.problem_.recompute_considered_[rec]:
                        if each_layer_per_recompute[layer][rec][inter][s] > 0:
                            act_mem += (each_layer_per_recompute[layer][rec][inter][s] *
                                        layer.memory_activation_rec_[rec] *
                                        activation_nums[inter][s])
        return act_mem


    def debug_print_manual_theoretical_memory(
            self,
            each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]],
            interleave_num: int = 1) -> None:
        """Log the per-stage theoretical memory implied by the solver model (debug aid)."""
        logger.info("%s Manual Theoretical Memory Analysis %s", "=" * 20, "=" * 20)

        if self.vpp_less_memory_:
            if self.seqpipe_:
                activation_nums = self.problem_.compute_activation_seq_nums(
                    self.num_of_stage_, interleave_num, self.seq_split_num_, self.num_of_micro_batch_, True)
            else:
                activation_nums = self.problem_.compute_less_activation_nums(
                    self.num_of_stage_, interleave_num)
        # Add if dual to decide whether dualpipe_v is used
        elif self.dual_:
            activation_nums = self.problem_.compute_activation_nums_dual(
                self.num_of_stage_, interleave_num, self.num_of_micro_batch_)
        else:
            if self.seqpipe_:
                activation_nums = self.problem_.compute_activation_seq_nums(
                    self.num_of_stage_, interleave_num, self.seq_split_num_, self.num_of_micro_batch_, False)
            else:
                activation_nums = self.problem_.compute_activation_nums(
                    self.num_of_stage_, interleave_num, self.num_of_micro_batch_)

        logger.info("Activation nums = %s", activation_nums)

        # compute for each stage
        for s in range(self.num_of_stage_):

            # parameter memory
            param_mem = self._compute_parameter_memory_manually_solver(each_layer_per_recompute, s, interleave_num)

            # head memory
            if s == 0:
                for head in self.layers_sorted_[Layer.type_enum.HEAD]:
                    if head.memory_parameter_ is not None:
                        param_mem += head.memory_parameter_

            # tail memory
            if s == self.num_of_stage_ - 1:
                for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
                    if tail.memory_parameter_ is not None:
                        param_mem += tail.memory_parameter_

            # act memory
            act_mem = self._calculate_activation_memory_solver(each_layer_per_recompute, s,
                                                               interleave_num, activation_nums)

            # overhead
            overhead = 0

            total = param_mem + act_mem + overhead + self.constant_memory_

            logger.info("Stage %d Manual Memory Analysis:", s)
            logger.info("Parameter Memory:     %.2f", param_mem)
            logger.info("Activation Memory:    %.2f", act_mem)
            logger.info("Memory Overhead:      %.2f", overhead)
            logger.info("Constant Memory:      %.2f", self.constant_memory_)
            logger.info("Total Theoretical Memory: %.2f", total)

    def split_layer_per_recompute(
            self,
            layer_per_recompute: Dict[Recompute.TYPE, List[List[int]]]
    ) -> Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]]:
        """Split aggregate per-recompute layer counts into counts per BODY layer."""
        each_layer_per_recompute = {}
        for layer in self.layers_sorted_[Layer.type_enum.BODY]:
            rest = layer.nb_layer_
            each_layer_per_recompute[layer] = {r: [] for r in Recompute.TYPE}
            for rec in Recompute.TYPE:
                for i in range(self.num_of_interleave_):
                    each_layer_per_recompute[layer][rec].append([0]*self.num_of_stage_)
                    for s in range(self.num_of_stage_):
                        subtract = min(layer_per_recompute[rec][i][s], rest)
                        layer_per_recompute[rec][i][s] -= subtract
                        rest -= subtract
                        each_layer_per_recompute[layer][rec][i][s] += subtract
        return each_layer_per_recompute

    def fuse_layer_per_recompute(
            self,
            each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]]
    ) -> Dict[Recompute.TYPE, List[List[int]]]:
        """Fuse per-layer recompute counts back into aggregate per-recompute-type totals."""
        all_layers_per_recompute = {r: [] for r in Recompute.TYPE}
        for rec in Recompute.TYPE:
            for i in range(self.num_of_interleave_):
                all_layers_per_recompute[rec].append([])
                for s in range(self.num_of_stage_):
                    all_layers_per_recompute[rec][i].append(sum(
                        each_layer_per_recompute[layer][rec][i][s]
                        for layer in self.layers_sorted_[Layer.type_enum.BODY]
                    ))
        return all_layers_per_recompute


    def simulate_manual(
            self,
            each_layer_per_recompute: Optional[Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]]] = None,
            show: bool = True,
            interleave_num: int = 1,

            interleave_num: int = 1,
            file_name: Optional[str] = None,
            sub_fig: Optional[plt.Figure] = None) -> float:
        """Run the simulator on a user-supplied per-layer recompute strategy."""
        logger.output("Simulating given strategy: %s", each_layer_per_recompute)

        for layer in self.layers_sorted_[Layer.type_enum.BODY]:
            for rec in Recompute.TYPE:
                if len(each_layer_per_recompute[layer][rec]) != interleave_num:
                    logger.error(
                        "For layer %s with recompute %s, %s does not match interleave number %s",
                        layer,
                        rec,
                        len(each_layer_per_recompute[layer][rec]),

                        rec,
                        len(each_layer_per_recompute[layer][rec]),
                        interleave_num,
                    )
                    return sys.maxsize

        for layer in self.layers_sorted_[Layer.type_enum.BODY]:
            for rec in Recompute.TYPE:
                if any(x < 0 for sublist in each_layer_per_recompute[layer][rec]
                    for x in sublist):
                    raise ValueError(
                        f"for {rec}, there is strategy less than 0 in "
                        f"{each_layer_per_recompute[layer][rec]}"
                    )

        forward_time = self.get_manual_fw_time(each_layer_per_recompute,
                                               interleave_num)
        recompute_overhead = self.get_manual_recompute_time(
            each_layer_per_recompute, interleave_num)
        stage_mem_par = 0
        stage_mem_act = 0
        if self.has_some_memory_info():
            stage_mem_par = self.get_manual_memory_parameter(
                each_layer_per_recompute, interleave_num=interleave_num)
            stage_mem_act = self.get_manual_memory_activation(
                each_layer_per_recompute, interleave_num=interleave_num)

        self.debug_print_manual_theoretical_memory(each_layer_per_recompute, interleave_num)

        return self.simulation(
            forward_time,
            recompute_overhead,
            stage_mem_par,
            stage_mem_act,

            file_name,
            sub_fig
        )

    def simulation(
            self,
            forward_time: List[List[float]],
            recompute_overhead: Union[int, List[List[float]]] = 0,
            stage_mem_par: Union[int, List[List[float]]] = 0,

            file_name: Optional[str] = None,
            sub_fig: Optional[plt.Figure] = None,
    ) -> float:
        """Run the low-level :class:`PipelineSimulator` and return its reported end time."""
        if self.has_some_memory_info():
            logger.output(
                "PipelineSimulator(\n\t%s, %s,"
                "\n\tblock_mem_act=%s,"
                "\n\tblock_mem_par=%s,"
                "\n\tlayer_recompute=%s,"

                recompute_overhead,
                self.vpp_less_memory_,
            )

            sim_method = "vpp2" if self.vpp_less_memory_ else "vpp"
            simulator = sim.PipelineSimulator(
                forward_time,
                self.num_of_micro_batch_,
                block_mem=stage_mem_act,
                block_mem_par=stage_mem_par,

                method=sim_method,
                sub_fig=sub_fig
            )
        else:
            logger.output(
                "PipelineSimulator(\n\t%s, %s,"
                "\n\tlayer_recompute=%s)"
                "\n\tless_memory=%s )",
                forward_time,

                self.num_of_micro_batch_,
                recompute_overhead,
                self.vpp_less_memory_,
            )
            simulator = sim.PipelineSimulator(
                forward_time,
                self.num_of_micro_batch_,
                layer_recompute=recompute_overhead,
                less_memory=self.vpp_less_memory_,