AIOB_simAI_workload_generator.py

"""
Copyright (c) 2021, Alibaba Group;
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
   http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import workload_generator.mocked_model.MockedDeepspeed
from workload_generator.mocked_model.MockedMegatron import *
from workload_generator.mocked_model.MockedModel import MockedParam, MockedModel
from utils.utils import CommType, get_params, get_comp_out, extract_averages
import os
from typing import List, Tuple
from collections import deque
import dataclasses
from enum import Enum

try:
    import torch
except ImportError as e:
    torch = None
    print("Failed to import 'torch'.")
import math
import re





@dataclasses.dataclass
class Work_Item:
    name: str = dataclasses.field(default="none")
    placeholder: int = dataclasses.field(default=-1)
    forward_compute_time: int = dataclasses.field(default=0)
    forward_comm: str = dataclasses.field(default="NONE")
    forward_comm_size: int = dataclasses.field(default=0)
    backward_compute_time: int = dataclasses.field(default=0)
    backward_comm: str = dataclasses.field(default="NONE")
    backward_comm_size: int = dataclasses.field(default=0)
    dp_compute_time: int = dataclasses.field(default=0)
    dp_comm: str = dataclasses.field(default="NONE")
    dp_comm_size: int = dataclasses.field(default=0)
    process_time: int = dataclasses.field(default=100)



def _get_aiob_compute_time(compute_cache, forward_or_backward, stage):
    compute_time_map = compute_cache
    if stage == "grad":
        prefix = stage + "_" + forward_or_backward
    elif stage == "embedding":
        prefix = "Emb"
    elif stage == "final":
        prefix = "attention" + "_" + forward_or_backward
    else:
        prefix = stage + "_" + forward_or_backward

    for key, value in compute_time_map.items():
        if prefix == key:

            compute_time = compute_time_map.get(key)
            return compute_time

    print("[warn] can't match any stage", stage)
    return 1


class LayerInfo:
    def __init__(self, layer_id, layer_name, param_count):
        self.layer_id = layer_id
        self.layer_name = layer_name
        self.param_count = param_count


class SIMAI_workload:
    def __init__(self, model, args, compute_cache=None):
        self.model = model
        self.args = args
        self.compute_cache = compute_cache
        self.workload = []
        self.seq_len = args.seq_length
        self.tp = args.tensor_model_parallel_size
        self.mbs = args.micro_batch
        if args.moe_enable:
            self.expert_model_parallel_size = args.expert_model_parallel_size
            self.num_experts = args.num_experts
            self.topk = args.moe_router_topk

    def get_model_details(self):
        layers = []
        visited = set()

        def traverse_model(model):
            if id(model) in visited:
                return
            visited.add(id(model))

            if self.args.enable_sequence_parallel:
                if (
                    isinstance(model, MegatronColumnLinear)
                    or isinstance(model, MegatronRowLinear)
                    or isinstance(model, MegatronEmbedding)
                    or isinstance(model, FusedLayernorm)
                ):
                    params = model.parameters()
                    param_count = sum(p.numel() for p in params)
                    layers.append(LayerInfo(model.layer_id, model.name, param_count))
                if isinstance(model, MOEMLP):
                    moe_params = model.parameters()
                    moe_param_count = sum(p.numel() for p in moe_params)
                    layers.append(LayerInfo(model.layer_id, model.name, moe_param_count))

            else:
                if (
                    isinstance(model, MegatronAttention)
                    or isinstance(model, MegatronMlp)
                    or isinstance(model, MegatronEmbedding)
                ):
                    params = model.parameters()
                    param_count = sum(p.numel() for p in params)
                    layers.append(LayerInfo(model.layer_id, model.name, param_count))

            for child in model.child_modules():
                traverse_model(child)

        traverse_model(model)

        return layers

    def _get_total_params(self):
        total_params = 0
        moe_param_count = 0
        layers = self.get_model_details()
        for layer in layers:
            total_params += layer.param_count
            if "moe" in layer.layer_name:
                moe_param_count += layer.param_count

        return total_params, moe_param_count

    def workload_generate_aiob(self):
        # args.world_size --> total gpus number
        self.ga_num = self.args.global_batch // (self.args.micro_batch * self.args.dp_num)
        if self.ga_num < 1:
            print(
                "[WARN]: ga num < 1, please confirm global_batch num and micro_batch num"
            )
        default_compute_time = 1
        compute_time = 0
        tp_comm_size = (
            2 * self.args.micro_batch * self.args.seq_length * self.args.hidden_size
        )
        layers = self.get_model_details()
        total_params, moe_param_count = self._get_total_params()
        # self.workload.append(Work_Item(name="norm", forward_compute_time=0,
        #                         forward_comm = "BROADCAST", forward_comm_size= 8*self.args.micro_batch*self.args.seq_length,
        #                         backward_compute_time=default_compute_time, backward_comm="NONE", backward_comm_size=0,
        #                         dp_compute_time=default_compute_time, dp_comm="NONE", dp_comm_size=0
        #                         ))
        forward_compute_time = _get_aiob_compute_time(
            self.compute_cache, "forward", "grad"
        )
        backward_compute_time = _get_aiob_compute_time(
            self.compute_cache, "backward", "grad"
        )
        self.workload.append(
            Work_Item(
                name="grad_gather",
                forward_compute_time=default_compute_time,
                forward_comm="NONE",
                forward_comm_size=0,
                backward_compute_time=default_compute_time,
                backward_comm="NONE",
                backward_comm_size=0,
                dp_compute_time=default_compute_time,
                dp_comm="ALLGATHER",
                dp_comm_size=2 * (total_params-moe_param_count),
            )
        )
        self.workload.append(
            Work_Item(
                name="grad_param_comm",
                forward_compute_time=default_compute_time,
                forward_comm="NONE",
                forward_comm_size=0,
                backward_compute_time=default_compute_time,
                backward_comm="NONE",
                backward_comm_size=0,
                dp_compute_time=default_compute_time,
                dp_comm="REDUCESCATTER",
                dp_comm_size=4 * (total_params-moe_param_count),
            )
        )
        self.workload.append(
            Work_Item(
                name="grad_param_compute",
                forward_compute_time=default_compute_time,
                forward_comm="NONE",
                forward_comm_size=0,
                backward_compute_time=forward_compute_time + backward_compute_time,
                backward_comm="NONE",
                backward_comm_size=0,
                dp_compute_time=default_compute_time,
                dp_comm="NONE",
                dp_comm_size=0,
            )
        )

        if not self.args.enable_sequence_parallel:
            self.workload.append(
                Work_Item(
                    name="layernorm",
                    forward_compute_time=default_compute_time,
                    forward_comm="NONE",
                    forward_comm_size=0,
                    backward_compute_time=default_compute_time,
                    backward_comm="ALLREDUCE",
                    backward_comm_size=2 * total_params,
                    dp_compute_time=default_compute_time,
                    dp_comm="NONE",
                    dp_comm_size=0,
                )
            )
        if args.tensor_model_parallel_size == 1 :
            emd_backward_comm = "NONE"
        else:
            emd_backward_comm = "ALLREDUCE"
        self.workload.append(
            Work_Item(
                name="embedding_grads",
                forward_compute_time=default_compute_time,
                forward_comm="NONE",
                forward_comm_size=0,
                backward_compute_time=default_compute_time,
                backward_comm=emd_backward_comm,
                backward_comm_size=tp_comm_size,
                dp_compute_time=default_compute_time,
                dp_comm="NONE",
                dp_comm_size=0,
            )
        )
        if self.args.expert_model_parallel_size != self.args.dp_num:
            self.workload.append(Work_Item(name="moe_grad_norm1", forward_compute_time=default_compute_time,
                                    forward_comm = "NONE", forward_comm_size= 0,
                                    backward_compute_time=default_compute_time, backward_comm="NONE", backward_comm_size=0,
                                    dp_compute_time=default_compute_time, dp_comm="ALLGATHER_DP_EP", dp_comm_size=2*moe_param_count
                                    ))
            self.workload.append(Work_Item(name="moe_grad_norm2", forward_compute_time=default_compute_time,
                                    forward_comm = "NONE", forward_comm_size= 0,
                                    backward_compute_time=default_compute_time, backward_comm="NONE", backward_comm_size=0,
                                    dp_compute_time=default_compute_time, dp_comm="REDUCESCATTER_DP_EP", dp_comm_size=4*moe_param_count
                                    ))
        for _ in range(self.ga_num):
            for layer in layers:
                name = layer.layer_name
                forward_comm = backward_comm = backward_comm_2 = "NONE"
                forward_comm_size = tp_comm_size
                emb_comm_size = tp_comm_size
                backward_comm_size = 0
                dp_comm = "NONE"
                dp_comm_size = 0
                if self.args.enable_sequence_parallel:
                    if "embedding" in name:
                        if args.tensor_model_parallel_size == 1 :
                            forward_comm = "NONE"
                            backward_comm = "NONE"
                        else:
                            forward_comm = "ALLREDUCE"
                            backward_comm = "NONE"
                        emb_compute_time = _get_aiob_compute_time(
                            self.compute_cache, "", "embedding"
                        )
                        self.workload.append(
                            Work_Item(
                                name=name,
                                forward_compute_time=emb_compute_time,
                                forward_comm=forward_comm,
                                forward_comm_size=emb_comm_size ,
                                backward_compute_time=default_compute_time,
                                backward_comm=backward_comm,
                                backward_comm_size=backward_comm_size,
                                dp_compute_time=backward_compute_time,
                                dp_comm=dp_comm,
                                dp_comm_size=dp_comm_size,
                            )
                        )
                    if "row" in name:
                        
                        forward_compute_time = _get_aiob_compute_time(
                        self.compute_cache, "forward", name.split("_")[0]
                        )
                        backward_compute_time = _get_aiob_compute_time(
                            self.compute_cache, "backward", name.split("_")[0]
                        )

                        if self.args.recompute_activations and 'attention' in name:
                            forward_compute_time *= 2
                        forward_compute_time = int(forward_compute_time / 2)
                        backward_compute_time = int(backward_compute_time / 2)
                        forward_comm_size_sp = tp_comm_size
                        if args.tensor_model_parallel_size == 1 :
                            forward_comm = "NONE"
                            backward_comm = "NONE"
                        else:
                            forward_comm = "REDUCESCATTER"
                            backward_comm = "ALLGATHER"
                        self.workload.append(
                                Work_Item(
                                    name=name,
                                    forward_compute_time=forward_compute_time,
                                    forward_comm=forward_comm,
                                    forward_comm_size=forward_comm_size,
                                    backward_compute_time=backward_compute_time,
                                    backward_comm=backward_comm,
                                    backward_comm_size=forward_comm_size_sp,#sp overlap allgather
                                    dp_compute_time=backward_compute_time,
                                    dp_comm=dp_comm,
                                    dp_comm_size=dp_comm_size,
                                )
                            )

                    elif "column" in name:
                        forward_compute_time = _get_aiob_compute_time(
                        self.compute_cache, "forward", name.split("_")[0]
                        )
                        backward_compute_time = _get_aiob_compute_time(
                            self.compute_cache, "backward", name.split("_")[0]
                        )

                        if self.args.recompute_activations and 'attention' in name:
                            forward_compute_time *= 2
                        forward_compute_time = int(forward_compute_time / 2)
                        backward_compute_time = int(backward_compute_time / 2)
                        if args.tensor_model_parallel_size == 1 :
                            forward_comm = "NONE"
                            backward_comm = "NONE"
                            backward_comm_2 = "NONE"
                        else:
                            forward_comm = "ALLGATHER"
                            backward_comm = "REDUCESCATTER"
                            backward_comm_2 = "ALLGATHER"
                        self.workload.append(
                                Work_Item(
                                    name=name,
                                    forward_compute_time=forward_compute_time,
                                    forward_comm=forward_comm,
                                    forward_comm_size=forward_comm_size,
                                    backward_compute_time=backward_compute_time,
                                    backward_comm=backward_comm,
                                    backward_comm_size=backward_comm_size,
                                    dp_compute_time=backward_compute_time,
                                    dp_comm=dp_comm,
                                    dp_comm_size=dp_comm_size,
                                )
                            )
                    elif "moelayer" in name:
                        forward_compute_time = _get_aiob_compute_time(
                        self.compute_cache, "forward", name.split("_")[0]
                        )
                        backward_compute_time = _get_aiob_compute_time(
                            self.compute_cache, "backward", name.split("_")[0]
                        )
                        if args.tensor_model_parallel_size == 1 :
                            forward_comm1 = "NONE"
                            forward_comm2 = "NONE"
                            forward_comm3 = "ALLTOALL_EP"
                            forward_comm4 = "NONE"
                            forward_comm5 = "NONE"
                            forward_comm6 = "ALLTOALL_EP"
                            forward_comm7 = "NONE"
                        else:
                            forward_comm1 = "ALLGATHER"
                            forward_comm2 = "ALLTOALL"
                            forward_comm3 = "ALLTOALL_EP"
                            forward_comm4 = "ALLGATHER"
                            forward_comm5 = "REDUCESCATTER"
                            forward_comm6 = "ALLTOALL_EP"
                            forward_comm7 = "ALLTOALL"
                        if args.expert_model_parallel_size != 1:
                            self.workload.append(Work_Item(name=name, forward_compute_time=forward_compute_time,
                                        forward_comm = forward_comm1, forward_comm_size= 2*self.mbs*self.seq_len*self.num_experts,
                                        backward_compute_time=backward_compute_time, backward_comm=forward_comm1, backward_comm_size=2*self.mbs*self.seq_len*self.num_experts,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm2, forward_comm_size= tp_comm_size//self.tp,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm2, backward_comm_size=tp_comm_size//self.tp,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm3, forward_comm_size= tp_comm_size*self.topk//self.tp,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm3, backward_comm_size=tp_comm_size*self.topk//self.tp,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm4, forward_comm_size= tp_comm_size*self.topk,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm5, backward_comm_size=tp_comm_size*self.topk,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm5, forward_comm_size= tp_comm_size*self.topk,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm4, backward_comm_size=tp_comm_size*self.topk,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm6, forward_comm_size= tp_comm_size*self.topk//self.tp,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm6, backward_comm_size=tp_comm_size*self.topk//self.tp,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm7, forward_comm_size= tp_comm_size//self.tp,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm7, backward_comm_size=tp_comm_size//self.tp,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                        else:
                            self.workload.append(Work_Item(name=name, forward_compute_time=forward_compute_time,
                                        forward_comm = forward_comm1, forward_comm_size= 2*self.mbs*self.seq_len*self.num_experts,
                                        backward_compute_time=backward_compute_time, backward_comm=forward_comm1, backward_comm_size=2*self.mbs*self.seq_len*self.num_experts,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm2, forward_comm_size= tp_comm_size//self.tp,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm2, backward_comm_size=tp_comm_size//self.tp,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm3, forward_comm_size=1,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm3, backward_comm_size=1,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm4, forward_comm_size= tp_comm_size*self.topk,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm4, backward_comm_size=tp_comm_size*self.topk,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm5, forward_comm_size= tp_comm_size*self.topk,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm4, backward_comm_size=tp_comm_size*self.topk,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm6, forward_comm_size=1,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm6, backward_comm_size=1,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                            self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                        forward_comm = forward_comm7, forward_comm_size= tp_comm_size//self.tp,
                                        backward_compute_time=default_compute_time, backward_comm=forward_comm7, backward_comm_size=tp_comm_size//self.tp,
                                        dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                        ))
                else:
                    if args.tensor_model_parallel_size == 1 :
                        forward_comm = "NONE"
                        backward_comm = "NONE"
                    else:

                        forward_comm = "ALLREDUCE"
                        backward_comm = "NONE"
                    if self.args.recompute_activations and 'attention' in name:
                        forward_compute_time *= 2
                    if "embedding" in name:
                        emb_compute_time = _get_aiob_compute_time(
                            self.compute_cache, "", "embedding"
                        )
                        self.workload.append(
                            Work_Item(
                                name=name,
                                forward_compute_time=emb_compute_time,
                                forward_comm=forward_comm,
                                forward_comm_size=forward_comm_size,
                                backward_compute_time=default_compute_time,
                                backward_comm=backward_comm,
                                backward_comm_size=backward_comm_size,
                                dp_compute_time=backward_compute_time,
                                dp_comm=dp_comm,
                                dp_comm_size=dp_comm_size,
                            )
                        )
                    else:
                        forward_compute_time = _get_aiob_compute_time(
                            self.compute_cache, "forward", name.split("_")[0]
                        )
                        backward_compute_time = _get_aiob_compute_time(
                            self.compute_cache, "backward", name.split("_")[0]
                        )
                        self.workload.append(
                            Work_Item(
                                name=name,
                                forward_compute_time=forward_compute_time,
                                forward_comm=forward_comm,
                                forward_comm_size=forward_comm_size,
                                backward_compute_time=backward_compute_time,
                                backward_comm=backward_comm,
                                backward_comm_size=backward_comm_size,
                                dp_compute_time=backward_compute_time,
                                dp_comm=dp_comm,
                                dp_comm_size=dp_comm_size,
                            )
                        )
            # compute_time = _get_aiob_compute_time(self.compute_cache, "forward", "embedding")
            # self.workload.append(Work_Item(name="embedding_norm", forward_compute_time=compute_time,
            #                         forward_comm = "ALLREDUCE", forward_comm_size= self.args.vocab_size*self.args.hidden_size*2,
            #                         backward_compute_time=default_compute_time, backward_comm="NONE", backward_comm_size=0,
            #                         dp_compute_time=default_compute_time, dp_comm="NONE", dp_comm_size=0
            #                         ))
        for i in range(3):
            self.workload.append(
                Work_Item(
                    name="cross_entropy" + str(i + 1),
                    forward_compute_time=compute_time,
                    forward_comm="ALLREDUCE",
                    forward_comm_size=self.args.seq_length * self.args.micro_batch * 4,
                    backward_compute_time=compute_time,
                    backward_comm="NONE",
                    backward_comm_size=0,
                    dp_compute_time=compute_time,
                    dp_comm="NONE",
                    dp_comm_size=0,
                )
            )

        for i in range(4):
            self.workload.append(
                Work_Item(
                    name="optimizer" + str(i + 1),
                    forward_compute_time=compute_time,
                    forward_comm="ALLREDUCE",
                    forward_comm_size=4,
                    backward_compute_time=compute_time,
                    backward_comm="NONE",
                    backward_comm_size=0,
                    dp_compute_time=compute_time,
                    dp_comm="NONE",
                    dp_comm_size=0,
                )
            )

    def workload_generate(self):
        # args.world_size --> total gpus number
        self.ga_num = self.args.global_batch // (self.args.micro_batch * self.args.dp_num)
        if self.ga_num < 1:
            print(
                "[WARN]: ga num < 1, please confirm global_batch num and micro_batch num"
            )
        default_compute_time = 1
        compute_time = 0
        tp_comm_size = (
            2 * self.args.micro_batch * self.args.seq_length * self.args.hidden_size
        )
        layers = self.get_model_details()
        total_params, moe_param_count = self._get_total_params()
        # print(f"Total params is {total_params}, moe params is {moe_param_count}")
        # self.workload.append(Work_Item(name="norm", forward_compute_time=0,
        #                         forward_comm = "BROADCAST", forward_comm_size= 8*self.args.micro_batch*self.args.seq_length,
        #                         backward_compute_time=default_compute_time, backward_comm="NONE", backward_comm_size=0,
        #                         dp_compute_time=default_compute_time, dp_comm="NONE", dp_comm_size=0
        #                         ))
        forward_compute_time = default_compute_time
        backward_compute_time = default_compute_time
        self.workload.append(
            Work_Item(
                name="grad_norm",
                forward_compute_time=forward_compute_time,
                forward_comm="ALLGATHER",
                forward_comm_size=2 * total_params,
                backward_compute_time=backward_compute_time,
                backward_comm="NONE",
                backward_comm_size=0,
                dp_compute_time=default_compute_time,
                dp_comm="REDUCESCATTER",
                dp_comm_size=4 * total_params,
            )
        )
        if not self.args.enable_sequence_parallel:
            self.workload.append(
                Work_Item(
                    name="layernorm",
                    forward_compute_time=default_compute_time,
                    forward_comm="NONE",
                    forward_comm_size=0,
                    backward_compute_time=default_compute_time,
                    backward_comm="ALLREDUCE",
                    backward_comm_size=2 * total_params,
                    dp_compute_time=default_compute_time,
                    dp_comm="NONE",
                    dp_comm_size=0,
                )
            )
        if args.expert_model_parallel_size != args.dp_num:
            self.workload.append(Work_Item(name="moe_grad_norm1", forward_compute_time=default_compute_time,
                                    forward_comm = "NONE", forward_comm_size= 0,
                                    backward_compute_time=default_compute_time, backward_comm="NONE", backward_comm_size=0,
                                    dp_compute_time=default_compute_time, dp_comm="ALLGATHER_DP_EP", dp_comm_size=2*moe_param_count
                                    ))
            self.workload.append(Work_Item(name="moe_grad_norm2", forward_compute_time=default_compute_time,
                                    forward_comm = "NONE", forward_comm_size= 0,
                                    backward_compute_time=default_compute_time, backward_comm="NONE", backward_comm_size=0,
                                    dp_compute_time=default_compute_time, dp_comm="REDUCESCATTER_DP_EP", dp_comm_size=4*moe_param_count
                                    ))
        for _ in range(self.ga_num):
            for layer in layers:
                name = layer.layer_name
                forward_comm = backward_comm = backward_comm_2 = "NONE"
                forward_comm_size = tp_comm_size
                backward_comm_size = tp_comm_size
                dp_comm = "NONE"
                dp_comm_size = 0
                if self.args.enable_sequence_parallel:
                    if "embedding" in name:
                        self.workload.append(
                            Work_Item(
                                name=name,
                                forward_compute_time=default_compute_time,
                                forward_comm=forward_comm,
                                forward_comm_size=forward_comm_size,
                                backward_compute_time=default_compute_time,
                                backward_comm=backward_comm,
                                backward_comm_size=backward_comm_size,
                                dp_compute_time=backward_compute_time,
                                dp_comm=dp_comm,
                                dp_comm_size=dp_comm_size,
                            )
                        )

                    if "row" in name:
                        if self.args.recompute_activations and 'attention' in name:
                            forward_comm_size *= 2
                        forward_comm = "REDUCESCATTER"
                        backward_comm = "ALLGATHER"
                        self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                    forward_comm = forward_comm, forward_comm_size= forward_comm_size,
                                    backward_compute_time=default_compute_time, backward_comm="NONE", backward_comm_size=tp_comm_size,
                                    dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                    ))
                    if "column" in name:
                        if self.args.recompute_activations and 'attention' in name:
                            forward_comm_size *= 2
                        forward_comm = "ALLGATHER"
                        forward_comm2 = "NONE"
                        backward_comm = "REDUCESCATTER"
                        backward_comm_2 = "ALLGATHER"
                        self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                    forward_comm = forward_comm, forward_comm_size= forward_comm_size,
                                    backward_compute_time=default_compute_time, backward_comm="NONE", backward_comm_size=0,
                                    dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                    ))
                    if "moelayer" in name:
                        forward_comm1 = "ALLGATHER"
                        forward_comm2 = "ALLTOALL"
                        forward_comm3 = "ALLTOALL_EP"
                        forward_comm4 = "ALLGATHER"
                        forward_comm5 = "REDUCESCATTER"
                        forward_comm6 = "ALLTOALL_EP"
                        forward_comm7 = "ALLTOALL"
                        self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                    forward_comm = forward_comm1, forward_comm_size= 2*self.seq_len*self.num_experts,
                                    backward_compute_time=default_compute_time, backward_comm=forward_comm1, backward_comm_size=2*self.seq_len*self.num_experts,
                                    dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                    ))
                        self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                    forward_comm = forward_comm2, forward_comm_size= tp_comm_size//self.tp,
                                    backward_compute_time=default_compute_time, backward_comm=forward_comm2, backward_comm_size=tp_comm_size//self.tp,
                                    dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                    ))
                        self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                    forward_comm = forward_comm3, forward_comm_size= tp_comm_size*self.topk//self.tp,
                                    backward_compute_time=default_compute_time, backward_comm=forward_comm3, backward_comm_size=tp_comm_size*self.topk//self.tp,
                                    dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                    ))
                        self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                    forward_comm = forward_comm4, forward_comm_size= tp_comm_size*self.topk//self.tp,
                                    backward_compute_time=default_compute_time, backward_comm=forward_comm4, backward_comm_size=tp_comm_size*self.topk,
                                    dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                    ))
                        self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                    forward_comm = forward_comm5, forward_comm_size= tp_comm_size*self.topk//self.tp,
                                    backward_compute_time=default_compute_time, backward_comm=forward_comm4, backward_comm_size=tp_comm_size*self.topk//self.tp,
                                    dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                    ))
                        self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                    forward_comm = forward_comm6, forward_comm_size= tp_comm_size*self.topk//self.tp,
                                    backward_compute_time=default_compute_time, backward_comm=forward_comm6, backward_comm_size=tp_comm_size*self.topk//self.tp,
                                    dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                    ))
                        self.workload.append(Work_Item(name=name, forward_compute_time=default_compute_time,
                                    forward_comm = forward_comm7, forward_comm_size= tp_comm_size//self.tp,
                                    backward_compute_time=default_compute_time, backward_comm=forward_comm7, backward_comm_size=tp_comm_size//self.tp,
                                    dp_compute_time=default_compute_time, dp_comm=dp_comm, dp_comm_size=dp_comm_size
                                    ))
                        
                else:
                    forward_comm = "ALLREDUCE"
                    backward_comm = "ALLREDUCE"
                    if self.args.recompute_activations and 'attention' in name:
                        forward_comm_size *= 2
                    if "embedding" in name:
                        self.workload.append(
                            Work_Item(
                                name=name,
                                forward_compute_time=default_compute_time,
                                forward_comm=forward_comm,
                                forward_comm_size=forward_comm_size,
                                backward_compute_time=default_compute_time,
                                backward_comm=backward_comm,
                                backward_comm_size=backward_comm_size,
                                dp_compute_time=backward_compute_time,
                                dp_comm=dp_comm,
                                dp_comm_size=dp_comm_size,
                            )
                        )
                    else:
                        self.workload.append(
                            Work_Item(
                                name=name,
                                forward_compute_time=default_compute_time,
                                forward_comm=forward_comm,
                                forward_comm_size=forward_comm_size,
                                backward_compute_time=default_compute_time,
                                backward_comm=backward_comm,
                                backward_comm_size=backward_comm_size,
                                dp_compute_time=default_compute_time,
                                dp_comm=dp_comm,
                                dp_comm_size=dp_comm_size,
                            )
                        )
            self.workload.append(
                Work_Item(
                    name="embedding_norm",
                    forward_compute_time=default_compute_time,
                    forward_comm="ALLREDUCE",
                    forward_comm_size=self.args.vocab_size * self.args.hidden_size * 2,
                    backward_compute_time=default_compute_time,
                    backward_comm="NONE",
                    backward_comm_size=0,
                    dp_compute_time=default_compute_time,
                    dp_comm="NONE",
                    dp_comm_size=0,
                )
            )
        for i in range(3):
            self.workload.append(
                Work_Item(
                    name="cross_entropy" + str(i + 1),
                    forward_compute_time=compute_time,
                    forward_comm="ALLREDUCE",
                    forward_comm_size=self.args.seq_length * self.args.micro_batch * 4,
                    backward_compute_time=compute_time,
                    backward_comm="NONE",
                    backward_comm_size=0,
                    dp_compute_time=compute_time,
                    dp_comm="NONE",
                    dp_comm_size=0,
                )
            )

        for i in range(4):
            self.workload.append(
                Work_Item(
                    name="optimizer" + str(i + 1),
                    forward_compute_time=compute_time,
                    forward_comm="ALLREDUCE",
                    forward_comm_size=4,
                    backward_compute_time=compute_time,
                    backward_comm="NONE",
                    backward_comm_size=0,
                    dp_compute_time=compute_time,
                    dp_comm="NONE",
                    dp_comm_size=0,
                )
            )

    def dump_file(self, filename):
        filename = filename + ".txt"

        pp_comm_value = 2 * self.args.micro_batch * self.args.seq_length * self.args.hidden_size
        if self.args.enable_sequence_parallel:
            pp_comm_value /= self.args.tensor_model_parallel_size

        pp_comm = (
            f"pp_comm: {pp_comm_value}"
            if self.args.pipeline_model_parallel != 1
            else "pp_comm: 0"
        )
        with open(filename, "w") as f:
            f.write((
                f"HYBRID_TRANSFORMER_FWD_IN_BCKWD model_parallel_NPU_group: {self.args.tensor_model_parallel_size} "
                f"ep: {self.args.expert_model_parallel_size} "
                f"pp: {self.args.pipeline_model_parallel} "
                f"vpp: {self.args.num_layers} "
                f"ga: {self.ga_num} all_gpus: {self.args.world_size} "
                f"checkpoints: 0 checkpoint_initiates: 0 "
            ) + pp_comm + "\n")

            f.write(str(len(self.workload)) + "\n")
            for item in self.workload:
                f.write(
                    "\t".join([str(getattr(item, k)) for k in item.__dict__.keys()])
                    + "\n"
                )


class simAI_MicroTest:
    def __init__(self, args):
        self.args = args
        self.workload = []

    def _simAI_microtest_convert(self, comm_type):
        if comm_type == "all_reduce" or comm_type == "allreduce":
            return "ALLREDUCE"
        elif comm_type == "all_gather" or comm_type == "allgather":
            return "ALLGATHER"
        elif comm_type == "reduce_scatter" or comm_type == "reducescatter":
            return "REDUCESCATTER"
        elif comm_type == "all_to_all" or comm_type == "alltoall":
            return "ALLTOALL"
        else:
            return

    def workload_generator(self):
        curr_size = self.args.begin_size
        default_compute_time = 1
        while curr_size <= self.args.end_size:
            self.workload.append(
                Work_Item(
                    name="micro_test",
                    forward_compute_time=default_compute_time,
                    forward_comm="NONE",
                    forward_comm_size=0,
                    backward_compute_time=default_compute_time,
                    backward_comm="NONE",
                    backward_comm_size=0,
                    dp_compute_time=default_compute_time,
                    dp_comm=self._simAI_microtest_convert(self.args.test_comm),
                    dp_comm_size=curr_size,
                    process_time=1,
                )
            )
            curr_size *= 2

    def dump_file(self, filename):
        filename = filename + ".txt"
        with open(filename, "w") as f:
            if not self.args.multi_all_reduce_enable:
                f.write(f"MICRO" + "\n")
                f.write(str(len(self.workload)) + "\n")
                for item in self.workload:
                    f.write(
                        "\t".join([str(getattr(item, k)) for k in item.__dict__.keys()])
                        + "\n"
                    )
            else:
                f.write(
                    f"HYBRID_TRANSFORMER_FWD_IN_BCKWD	model_parallel_NPU_group: {self.args.tensor_model_parallel_size} \
                        expert_parallel_npu_group: {self.args.expert_model_parallel_size} pp: {self.args.pipeline_model_parallel} \
                        ga: {self.ga_num} all_gpus: {self.args.world_size} checkpoints: 0 checkpoint_initiates: 0"
                    + "\n"
                )
                f.write(str(len(self.workload)) + "\n")
                for item in self.workload:
                    f.write(
                        "\t".join([str(getattr(item, k)) for k in item.__dict__.keys()])
                        + "\n"
                    )


if __name__ == "__main__":
    args = get_params()
    print(args)
    model = MegatronModel(args)
    result_dir = "results/workload/"
    if not os.path.isdir(result_dir):
        os.makedirs(result_dir)
    filename = f"{args.gpu_type}-{args.model_name}-world_size{args.world_size}-tp{args.tensor_model_parallel_size}-pp{args.pipeline_model_parallel}-ep{args.expert_model_parallel_size}-gbs{args.global_batch}-mbs{args.micro_batch}-seq{args.seq_length}-MOE-{args.moe_enable}-GEMM-{args.moe_grouped_gemm}-flash_attn-{args.use_flash_attn}"
    filepath = os.path.join(result_dir, filename)
    params = model.parameters()
    # work = SIMAI_workload(model, args, GPU_Tensor_core.A100, "gpt13B")
    # name_layers = work.workload_generate()
    # work.dump_file("test")
    print(sum(p.numel() for p in params))
    if args.aiob_enable:
        params = model.parameters()
        args.model_param = sum(p.numel() for p in params)
        if args.comp_filepath == None:

            comp_filepath = get_comp_out(args)

            compute_cache = extract_averages(comp_filepath,args)
        else:
            print("comp_filepath:", args.comp_filepath)
            comp_filepath = args.comp_filepath
            compute_cache = extract_averages(comp_filepath,args)

        print("compute_cache = {")
        for key, value in compute_cache.items():
            print(f"    '{key}' : {value},")
        print("}")
        work = SIMAI_workload(
            model, args,compute_cache
        )
        name_layers = work.workload_generate_aiob()

        work.dump_file(filepath)
        print("workload save in :", filepath)
    # print(args)
    else:

        work = SIMAI_workload(model, args, None)
        name_layers = work.workload_generate()
        work.dump_file(filepath)
        print(f"workload save in : {filepath}.txt")