methods/elasticdnn/hugging_face/internal_adapter.py

from ..api.model import ElasticDNN_OfflineFMModel, ElasticDNN_OfflineMDModel
from .user_impl import HuggingFaceModelAPI

from typing import List
from data.dataloader import build_dataloader
# from methods.elasticdnn.api.online_model import ElasticDNN_OnlineModel
from methods.elasticdnn.api.online_model_v2 import ElasticDNN_OnlineModel

import torch
import sys
from torch import nn
from methods.elasticdnn.api.model import ElasticDNN_OfflineSegFMModel, ElasticDNN_OfflineSegMDModel
from methods.elasticdnn.api.algs.md_pretraining_wo_fbs import ElasticDNN_MDPretrainingWoFBSAlg
from methods.elasticdnn.model.base import ElasticDNNUtil
from methods.elasticdnn.pipeline.offline.fm_to_md.base import FM_to_MD_Util
from methods.elasticdnn.pipeline.offline.fm_to_md.vit import FM_to_MD_ViT_Util
from methods.elasticdnn.pipeline.offline.fm_lora.base import FMLoRA_Util
from methods.elasticdnn.pipeline.offline.fm_lora.vit import FMLoRA_ViT_Util
from methods.elasticdnn.model.bert import ElasticBertUtil
from utils.common.file import ensure_dir
from utils.dl.common.model import LayerActivation, get_module, get_parameter, get_super_module
from utils.common.exp import save_models_dict_for_init, get_res_save_dir
from data import build_scenario
from utils.dl.common.loss import CrossEntropyLossSoft
import torch.nn.functional as F
from utils.dl.common.env import create_tbwriter
import os
from utils.common.log import logger
from utils.common.data_record import write_json
# from methods.shot.shot import OnlineShotModel
from methods.feat_align.main import OnlineFeatAlignModel, FeatAlignAlg
import tqdm
from methods.feat_align.mmd import mmd_rbf
from copy import deepcopy
from typing import Optional, Union
import torch
from torch import nn 
from einops import rearrange, repeat
from einops.layers.torch import Rearrange
import tqdm
from methods.elasticdnn.model.vit import Linear_WrappedWithFBS

from utils.dl.common.model import get_model_device, get_model_size, set_module, get_module
import torch
from abc import abstractmethod


class ElasticDNN_OfflineFMModel_for_HuggingFaceFM(ElasticDNN_OfflineFMModel):
    def set_hugging_face_api(self, hugging_face_api: HuggingFaceModelAPI):
        self.hugging_face_api = hugging_face_api
        
    def get_accuracy(self, test_loader, *args, **kwargs):
        return self.hugging_face_api.get_accuracy(self.models_dict['main'], test_loader, self.device, *args, **kwargs)
    
    def infer(self, x, *args, **kwargs):
        return self.hugging_face_api.infer(self.models_dict['main'], x, *args, **kwargs)
    
    def get_required_model_components(self) -> List[str]:
        return ['main']
    
    def generate_md_by_reducing_width(self, reducing_width_ratio, samples: torch.Tensor):
        res = FM_to_MD_HuggingFaceFM_Util()
        res.set_hugging_face_api(self.hugging_face_api)
        return res.init_md_from_fm_by_reducing_width_with_perf_test(self.models_dict['main'], reducing_width_ratio, samples)
    
    def forward_to_get_task_loss(self, x, y, *args, **kwargs):
        return self.hugging_face_api.forward_to_get_task_loss(self.models_dict['main'], x, y)
    
    def get_feature_hook(self) -> LayerActivation:
        return self.hugging_face_api.get_feature_hook(self.models_dict['main'], self.device)
        
    def get_elastic_dnn_util(self) -> ElasticDNNUtil:
        res = ElasticHuggingFaceFMUtil()
        res.set_hugging_face_api(self.hugging_face_api)
        return res
    
    def get_lora_util(self) -> FMLoRA_Util:
        res = FMLoRA_HuggingFaceFM_Util()
        res.set_hugging_face_api(self.hugging_face_api)
        return res
    
    def get_task_head_params(self):
        return self.hugging_face_api.get_task_head_params(self.models_dict['main'])
    
    
class ElasticDNN_OfflineMDModel_for_HuggingFaceFM(ElasticDNN_OfflineMDModel):
    def set_hugging_face_api(self, hugging_face_api: HuggingFaceModelAPI):
        self.hugging_face_api = hugging_face_api
        
    def get_required_model_components(self) -> List[str]:
        return ['main']
    
    def get_accuracy(self, test_loader, *args, **kwargs):
        return self.hugging_face_api.get_accuracy(self.models_dict['main'], test_loader, self.device, *args, **kwargs)
    
    def infer(self, x, *args, **kwargs):
        return self.hugging_face_api.infer(self.models_dict['main'], x, *args, **kwargs)
    
    def forward_to_get_task_loss(self, x, y, *args, **kwargs):
        return self.hugging_face_api.forward_to_get_task_loss(self.models_dict['main'], x, y)
    
    def get_feature_hook(self) -> LayerActivation:
        return self.hugging_face_api.get_feature_hook(self.models_dict['main'], self.device)
    
    def get_distill_loss(self, student_output, teacher_output):
        return CrossEntropyLossSoft()(student_output, teacher_output)
    
    def get_matched_param_of_fm(self, self_param_name, fm: nn.Module):
        if any([k in self_param_name for k in ['fbs', 'cls_token', 'pos_embed']]):
            return None
        
        p = get_parameter(self.models_dict['main'], self_param_name)
        if p.dim() == 0:
            return None
        elif p.dim() == 1:
            return None
        
        layers_name = self.hugging_face_api.get_qkv_proj_ff1_ff2_layer_names()
        if len(layers_name[0]) == 4:
            
            
            qkv_names = [layer[0] for layer in layers_name]
            qkv_proj_names = [layer[1] for layer in layers_name]
            ff1_names = [layer[-2] for layer in layers_name]
            ff2_names = [layer[-1] for layer in layers_name]
            
            qkv_weight_names = [n + '.weight' for n in qkv_names]
        
            if self_param_name in qkv_weight_names:
                ss = self_param_name.split('.')
                
                fm_qkv_name = '.'.join(ss[0: -1]) + '.qkv'
                fm_qkv = get_module(fm, fm_qkv_name)
                
                fm_abs_name = '.'.join(ss[0: -1]) + '.abs'
                fm_abs = get_module(fm, fm_abs_name)
                
                # print(fm_qkv_name, fm_abs_name, fm)
                
                return torch.cat([
                    fm_qkv.weight.data, # task-agnositc params
                    torch.cat([(_abs[0].weight.T @ _abs[1].weight.T).T for _abs in fm_abs], dim=0) # task-specific params (LoRA)
                ], dim=0)
        else:
            q_names = [layer[0] for layer in layers_name]
            k_names = [layer[1] for layer in layers_name]
            v_names = [layer[2] for layer in layers_name]
            qkv_proj_names = [layer[3] for layer in layers_name]
            ff1_names = [layer[-2] for layer in layers_name]
            ff2_names = [layer[-1] for layer in layers_name]
            
            qkv_weight_names = [n + '.weight' for n in q_names + k_names + v_names]
        
            if self_param_name in qkv_weight_names:
            
                ss = self_param_name.split('.')
                # raise NotImplementedError() # TODO:
                fm_qkv_name = '.'.join(ss[0: -1]) + '.qkv'
                fm_qkv = get_module(fm, fm_qkv_name)
                
                fm_abs_name = '.'.join(ss[0: -1]) + '.ab'
                fm_abs = get_module(fm, fm_abs_name)
                
                # print(fm_qkv_name, fm_abs_name, fm)
                
                return torch.cat([
                    fm_qkv.weight.data, # task-agnositc params
                    fm_abs[1].weight @ fm_abs[0].weight
                ], dim=0)
            
        # elif 'to_qkv.bias' in self_param_name:
        #     ss = self_param_name.split('.')
            
        #     fm_qkv_name = '.'.join(ss[0: -2]) + '.qkv.bias'
        #     return get_parameter(fm, fm_qkv_name)
        
        ff1_weight_names = [n + '.linear.weight' for n in ff1_names]
        ff2_weight_names = [n + '.weight' for n in ff2_names]
            
        if self_param_name in ff1_weight_names:
            fm_param_name = self_param_name.replace('.linear', '')
            return get_parameter(fm, fm_param_name)

        if self_param_name in ff2_weight_names:
            fm_param_name = self_param_name
            return get_parameter(fm, fm_param_name)
        
        return None
    
    
class ElasticHuggingFaceFMUtil(ElasticDNNUtil):
    def set_hugging_face_api(self, hugging_face_api: HuggingFaceModelAPI):
        self.hugging_face_api = hugging_face_api
        
    def convert_raw_dnn_to_master_dnn(self, raw_dnn: nn.Module, r: float, ignore_layers=[]):
        assert len(ignore_layers) == 0, 'not supported yet'

        raw_vit = deepcopy(raw_dnn)
        
        # set_module(module, 'patch_embed.proj', ProjConv_WrappedWithFBS(module.patch_embed.proj, r))
        layers = self.hugging_face_api.get_qkv_proj_ff1_ff2_layer_names()
        ff1_names = [layer[-2] for layer in layers]
                
        for name, module in raw_vit.named_modules():
            # if name.endswith('attn'):
            #     set_module(module, 'qkv', ToQKV_WrappedWithFBS(module.qkv, r))
            if name in ff1_names:
                # set_module(get_super_module(module, name), name.split('.')[-1], Linear_WrappedWithFBS(module, r))
                set_module(raw_vit, name, Linear_WrappedWithFBS(module, r))
        
        return raw_vit
    
    def set_master_dnn_sparsity(self, master_dnn: nn.Module, sparsity: float):
        return super().set_master_dnn_sparsity(master_dnn, sparsity)
    
    def select_most_rep_sample(self, master_dnn: nn.Module, samples: torch.Tensor):
        # print(samples)
        # return samples[0].unsqueeze(0)
        res = {k: v[0: 1] for k, v in samples.items()}
        return res
        
    def extract_surrogate_dnn_via_samples(self, master_dnn: nn.Module, samples: torch.Tensor, return_detail=False):
        sample = self.select_most_rep_sample(master_dnn, samples)
        # assert sample.dim() == 4 and sample.size(0) == 1
        
        # print('before')
        master_dnn.eval()
        self.clear_cached_channel_attention_in_master_dnn(master_dnn)
        with torch.no_grad():
            master_dnn_output = master_dnn(**sample)
            
        # print('after')
        
        boosted_vit = deepcopy(master_dnn)
        
        def get_unpruned_indexes_from_channel_attn(channel_attn: torch.Tensor, k):
            assert channel_attn.size(0) == 1, 'use A representative sample to generate channel attentions'
            
            res = channel_attn[0].nonzero(as_tuple=True)[0] # should be one-dim
            return res
        
        unpruned_indexes_of_layers = {}
        
        layers_name = self.hugging_face_api.get_qkv_proj_ff1_ff2_layer_names()
        ff1_names = [layer[-2] for layer in layers]
        ff2_names = [layer[-1] for layer in layers]
        
        for ff1_name, ff2_name in zip(ff1_names, ff2_names):
            ff_0 = get_module(boosted_vit, ff1_name)
            # ff_0_unpruned_indexes = get_unpruned_indexes_from_channel_attn(ff_0.cached_channel_attention, k)
            ff_0_pruned_indexes = ff_0.k_takes_all.cached_i[0].sort()[0]
            ff_0_unpruned_indexes = torch.LongTensor([ii for ii in range(ff_0.cached_channel_attention.size(1)) if ii not in ff_0_pruned_indexes])
            new_ff_0 = nn.Linear(ff_0.linear.in_features, ff_0_unpruned_indexes.size(0), ff_0.linear.bias is not None)
            new_ff_0.weight.data.copy_(ff_0.linear.weight.data[ff_0_unpruned_indexes])
            if ff_0.linear.bias is not None:
                new_ff_0.bias.data.copy_(ff_0.linear.bias.data[ff_0_unpruned_indexes])
            # set_module(get_super_module(ff_0, ff1_name), ff1_name.split('.')[-1], 
            #            nn.Sequential(new_ff_0, StaticFBS(ff_0.cached_channel_attention[:, ff_0_unpruned_indexes])))
            set_module(boosted_vit, ff1_name, 
                       nn.Sequential(new_ff_0, StaticFBS(ff_0.cached_channel_attention[:, ff_0_unpruned_indexes])))
            
            ff_1 = get_module(boosted_vit, ff2_name)
            new_ff_1 = nn.Linear(ff_0_unpruned_indexes.size(0), ff_1.out_features, ff_1.bias is not None)
            new_ff_1.weight.data.copy_(ff_1.weight.data[:, ff_0_unpruned_indexes])
            if ff_1.bias is not None:
                new_ff_1.bias.data.copy_(ff_1.bias.data)
            # set_module(get_super_module(ff_1), ff2_name.split('.')[-1],  new_ff_1)
            set_module(boosted_vit, ff2_name, new_ff_1)
            
            unpruned_indexes_of_layers[f'{ff1_name}.0.weight'] = ff_0_unpruned_indexes
        
        surrogate_dnn = boosted_vit
        surrogate_dnn.eval()
        surrogate_dnn = surrogate_dnn.to(get_model_device(master_dnn))
        # logger.debug(surrogate_dnn)
        with torch.no_grad():
            surrogate_dnn_output = surrogate_dnn(**sample)
            
        output_diff = ((surrogate_dnn_output - master_dnn_output) ** 2).sum()
        # assert output_diff < 1e-4, output_diff
        logger.info(f'output diff of master and surrogate DNN: {output_diff}')
        logger.debug(f'example output of master/surrogate: {master_dnn_output.sum(0)[0: 10]}, {surrogate_dnn_output.sum(0)[0: 10]}')
        # logger.info(f'\nonly prune mlp!!!!\n')
        # logger.info(f'\nonly prune mlp!!!!\n')
        
        if return_detail:
            return boosted_vit, unpruned_indexes_of_layers
        
        return boosted_vit
    
    def extract_surrogate_dnn_via_samples_with_perf_test(self, master_dnn: nn.Module, samples: torch.Tensor, return_detail=False):
        master_dnn_size = get_model_size(master_dnn, True)
        master_dnn_latency = self._get_model_latency(master_dnn, samples, 50, 
                                               get_model_device(master_dnn), 50, False)
        
        res = self.extract_surrogate_dnn_via_samples(master_dnn, samples, return_detail)
        if not return_detail:
            surrogate_dnn = res
        else:
            surrogate_dnn, unpruned_indexes_of_layers = res
        surrogate_dnn_size = get_model_size(surrogate_dnn, True)
        surrogate_dnn_latency = self._get_model_latency(master_dnn, samples, 50, 
                                               get_model_device(master_dnn), 50, False)

        logger.info(f'master DNN ({master_dnn_size:.3f}MB, {master_dnn_latency:.4f}s/sample) -> '
                    f'surrogate DNN ({surrogate_dnn_size:.3f}MB, {surrogate_dnn_latency:.4f}s/sample)\n'
                    f'(model size: ↓ {(master_dnn_size / surrogate_dnn_size):.2f}x, '
                    f'latency: ↓ {(master_dnn_latency / surrogate_dnn_latency):.2f}x)')
        
        return res
    
    def _get_model_latency(self, model: torch.nn.Module, model_input_size, sample_num: int, 
                           device: str, warmup_sample_num: int, return_detail=False):
        import time
        
        if isinstance(model_input_size, tuple):
            dummy_input = torch.rand(model_input_size).to(device)
        else:
            dummy_input = model_input_size
            
        model = model.to(device)
        model.eval()
        
        # warm up
        with torch.no_grad():
            for _ in range(warmup_sample_num):
                model(**dummy_input)
                
        infer_time_list = []
                
        if device == 'cuda' or 'cuda' in str(device):
            with torch.no_grad():
                for _ in range(sample_num):
                    s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
                    s.record()
                    model(**dummy_input)
                    e.record()
                    torch.cuda.synchronize()
                    cur_model_infer_time = s.elapsed_time(e) / 1000.
                    infer_time_list += [cur_model_infer_time]

        else:
            with torch.no_grad():
                for _ in range(sample_num):
                    start = time.time()
                    model(**dummy_input)
                    cur_model_infer_time = time.time() - start
                    infer_time_list += [cur_model_infer_time]
                    
        avg_infer_time = sum(infer_time_list) / sample_num

        if return_detail:
            return avg_infer_time, infer_time_list
        return avg_infer_time
    
    
class FMLoRA_HuggingFaceFM_Util(FMLoRA_Util):
    def set_hugging_face_api(self, hugging_face_api: HuggingFaceModelAPI):
        self.hugging_face_api = hugging_face_api
    
    @torch.no_grad()
    def add_lora_ab_to_fm(self, fm: nn.Module, ab_r: int, samples: dict):
        fm.eval()
        
        if isinstance(samples, dict):
            o1 = fm(**samples)
        else:
            o1 = fm(samples)
        
        layers_name = self.hugging_face_api.get_qkv_proj_ff1_ff2_layer_names()
        if len(layers_name[0]) == 4:
            qkv_names = [layer[0] for layer in layers_name]
            
            from ..pipeline.offline.fm_lora.vit import ToQKV_WrappedWithLoRA
            for name, module in fm.named_modules():
                if name in qkv_names:
                    set_module(fm, name, ToQKV_WrappedWithLoRA(module, ab_r))
        else:
            qkv_names = [layer[0] for layer in layers_name] + [layer[1] for layer in layers_name] + [layer[2] for layer in layers_name]
            
            from ..pipeline.offline.fm_lora.bert import ToQKV_WrappedWithLoRA
            for name, module in fm.named_modules():
                if name in qkv_names:
                    set_module(fm, name, ToQKV_WrappedWithLoRA(module, ab_r))
        
        
        
        if isinstance(samples, dict):
            o2 = fm(**samples)
        else:
            o2 = fm(samples)
        
        if isinstance(o1, tuple):
            o1 = o1[-1]
            o2 = o2[-1]
        output_diff = ((o1 - o2) ** 2).sum()
        assert output_diff < 1e-5
        
        return fm
    
    @torch.no_grad()
    def absorb_lora_and_recover_net_structure(self, fm: nn.Module, samples: dict):       
        fm.eval()
        # print('absorb lora before')
        if isinstance(samples, dict):
            o1 = fm(**samples)
        else:
            o1 = fm(samples)
        
        from ..pipeline.offline.fm_lora.vit import ToQKV_WrappedWithLoRA as ToQKV_WrappedWithLoRA1
        from ..pipeline.offline.fm_lora.bert import ToQKV_WrappedWithLoRA as ToQKV_WrappedWithLoRA2
        
        for name, module in fm.named_modules():
            if isinstance(module, ToQKV_WrappedWithLoRA1):
            
                qkv = module.qkv
                fm_abs = module.abs

                fm_abs_weight = torch.cat([_abs[1].weight @ _abs[0].weight for _abs in fm_abs], dim=0)
                qkv.weight.add_(fm_abs_weight)
                
                set_module(fm, name, qkv)
            
            elif isinstance(module, ToQKV_WrappedWithLoRA2):
            
                fc = module.fc
                ab = module.ab

                fc.weight.add_(ab[1].weight @ ab[0].weight)
                
                set_module(fm, name, fc)
        
        # print('absorb lora after')
        if isinstance(samples, dict):
            o2 = fm(**samples)
        else:
            o2 = fm(samples)
        
        if isinstance(o1, tuple):
            o1 = o1[-1]
            o2 = o2[-1]
        output_diff = ((o1 - o2) ** 2).sum()
        assert output_diff < 1e-6, output_diff
        
        return fm
        
        
class FM_to_MD_HuggingFaceFM_Util(FM_to_MD_Util):
    def set_hugging_face_api(self, hugging_face_api: HuggingFaceModelAPI):
        self.hugging_face_api = hugging_face_api
        
    def init_md_from_fm_by_reducing_width(self, fm: nn.Module, reducing_width_ratio: int) -> nn.Module:
        fm_vit = deepcopy(fm)
        
        # for block in fm_vit.bert.encoder.layer:
        #     set_module(block, 'attention.self', BertSelfAttentionPrunable.init_from_exist_self_attn(block.attention.self))
        
        def _f(n):
            return int(n // reducing_width_ratio)
        
        # def _rand_indexes(n):
            # return torch.randperm(n)[0: int(n // reducing_width_ratio)]
            
        def l1_max_indexes(p: torch.Tensor, dim=0):
            assert dim in [0, 1]
            assert p.dim() in [1, 2, 4]
            
            if dim == 1:
                p = p.T
            
            p_norm = p.abs().contiguous().view(p.size(0), -1).sum(dim=1)
            n = p.size(0)
            return p_norm.argsort(descending=True)[0: int(n // reducing_width_ratio)].sort()[0]
        
        layers_name = self.hugging_face_api.get_qkv_proj_ff1_ff2_layer_names()
        if len(layers_name[0]) == 6:
            q_names = [layer[0] for layer in layers_name]
            k_names = [layer[1] for layer in layers_name]
            v_names = [layer[2] for layer in layers_name]
            qkv_proj_names = [layer[3] for layer in layers_name]
            ff1_names = [layer[-2] for layer in layers_name]
            ff2_names = [layer[-1] for layer in layers_name]
            
            for q_name, k_name, v_name, qkv_proj_name, ff1_name, ff2_name in zip(q_names, k_names, v_names, qkv_proj_names, ff1_names, ff2_names):
                for k in [q_name, k_name, v_name]:
                    qkv = get_module(fm_vit, k)

                    new_qkv = nn.Linear(qkv.in_features, _f(qkv.out_features), 
                                        qkv.bias is not None, qkv.weight.device)
                    indexes = l1_max_indexes(qkv.weight.data, 0)
                    
                    new_qkv.weight.data.copy_(qkv.weight.data[indexes])
                    if qkv.bias is not None:
                        new_qkv.bias.data.copy_(qkv.bias.data[indexes])
                    set_module(fm_vit, k, new_qkv)
                    
                proj = get_module(fm_vit, qkv_proj_name)
                new_proj = nn.Linear(_f(proj.in_features), proj.out_features, 
                                    proj.bias is not None, proj.weight.device)
                new_proj.weight.data.copy_(proj.weight.data[:, l1_max_indexes(proj.weight.data, 1)])
                if proj.bias is not None:
                    new_proj.bias.data.copy_(proj.bias.data)
                set_module(fm_vit, qkv_proj_name, new_proj)
                
                fc1 = get_module(fm_vit, ff1_name)
                new_fc1 = nn.Linear(fc1.in_features, _f(fc1.out_features), 
                                    fc1.bias is not None, fc1.weight.device)
                indexes = l1_max_indexes(fc1.weight.data, 0)
                new_fc1.weight.data.copy_(fc1.weight.data[indexes])
                if fc1.bias is not None:
                    new_fc1.bias.data.copy_(fc1.bias.data[indexes])
                set_module(fm_vit, ff1_name, new_fc1)

                fc2 = get_module(fm_vit, ff2_name)
                new_fc2 = nn.Linear(_f(fc2.in_features), fc2.out_features, 
                                    fc2.bias is not None, fc2.weight.device)
                new_fc2.weight.data.copy_(fc2.weight.data[:, l1_max_indexes(fc2.weight.data, 1)])
                if fc2.bias is not None:
                    new_fc2.bias.data.copy_(fc2.bias.data)
                set_module(fm_vit, ff2_name, new_fc2)
                
        if len(layers_name[0]) == 4:
            qkv_names = [layer[0] for layer in layers_name]
            qkv_proj_names = [layer[1] for layer in layers_name]
            ff1_names = [layer[-2] for layer in layers_name]
            ff2_names = [layer[-1] for layer in layers_name]
            
            for qkv_name, qkv_proj_name, ff1_name, ff2_name in zip(qkv_names, qkv_proj_names, ff1_names, ff2_names):
                qkv = get_module(fm_vit, qkv_name)
                new_qkv = nn.Linear(qkv.in_features, _f(qkv.out_features), 
                                    qkv.bias is not None, qkv.weight.device)
                indexes = l1_max_indexes(qkv.weight.data, 0)
                
                new_qkv.weight.data.copy_(qkv.weight.data[indexes])
                if qkv.bias is not None:
                    new_qkv.bias.data.copy_(qkv.bias.data[indexes])
                set_module(fm_vit, qkv_name, new_qkv)
                    
                proj = get_module(fm_vit, qkv_proj_name)
                new_proj = nn.Linear(_f(proj.in_features), proj.out_features, 
                                    proj.bias is not None, proj.weight.device)
                new_proj.weight.data.copy_(proj.weight.data[:, l1_max_indexes(proj.weight.data, 1)])
                if proj.bias is not None:
                    new_proj.bias.data.copy_(proj.bias.data)
                set_module(fm_vit, qkv_proj_name, new_proj)
                
                fc1 = get_module(fm_vit, ff1_name)
                new_fc1 = nn.Linear(fc1.in_features, _f(fc1.out_features), 
                                    fc1.bias is not None, fc1.weight.device)
                indexes = l1_max_indexes(fc1.weight.data, 0)
                new_fc1.weight.data.copy_(fc1.weight.data[indexes])
                if fc1.bias is not None:
                    new_fc1.bias.data.copy_(fc1.bias.data[indexes])
                set_module(fm_vit, ff1_name, new_fc1)

                fc2 = get_module(fm_vit, ff2_name)
                new_fc2 = nn.Linear(_f(fc2.in_features), fc2.out_features, 
                                    fc2.bias is not None, fc2.weight.device)
                new_fc2.weight.data.copy_(fc2.weight.data[:, l1_max_indexes(fc2.weight.data, 1)])
                if fc2.bias is not None:
                    new_fc2.bias.data.copy_(fc2.bias.data)
                set_module(fm_vit, ff2_name, new_fc2)
            
        return fm_vit
    
    def init_md_from_fm_by_reducing_width_with_perf_test(self, fm: nn.Module, reducing_width_ratio: int,
                                                         samples: torch.Tensor) -> nn.Module:
        fm_size = get_model_size(fm, True)
        fm_latency = self._get_model_latency(fm, samples, 20, 
                                               get_model_device(fm), 20, False)
        
        master_dnn = self.init_md_from_fm_by_reducing_width(fm, reducing_width_ratio)
        master_dnn_size = get_model_size(master_dnn, True)
        logger.debug(f'inited master DNN: {master_dnn}')
        master_dnn_latency = self._get_model_latency(master_dnn, samples, 20, 
                                               get_model_device(master_dnn), 20, False)

        logger.info(f'init master DNN (w/o FBS yet) by reducing foundation model\'s width (by {reducing_width_ratio:d}x)')
        logger.info(f'foundation model ({fm_size:.3f}MB, {fm_latency:.4f}s/sample) -> '
                    f'master DNN ({master_dnn_size:.3f}MB, {master_dnn_latency:.4f}s/sample)\n'
                    f'(model size: ↓ {(fm_size / master_dnn_size):.2f}x, '
                    f'latency: ↓ {(fm_latency / master_dnn_latency):.2f}x)')
        
        return master_dnn
    
    def _get_model_latency(self, model: torch.nn.Module, model_input_size, sample_num: int, 
                           device: str, warmup_sample_num: int, return_detail=False):
        import time
        
        if isinstance(model_input_size, tuple):
            dummy_input = torch.rand(model_input_size).to(device)
        else:
            dummy_input = model_input_size
            
        model = model.to(device)
        model.eval()
        
        # warm up
        with torch.no_grad():
            for _ in range(warmup_sample_num):
                if isinstance(dummy_input, dict):
                    model(**dummy_input)
                else:
                    model(dummy_input)
                
        infer_time_list = []
                
        if device == 'cuda' or 'cuda' in str(device):
            with torch.no_grad():
                for _ in range(sample_num):
                    s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
                    s.record()
                    
                    if isinstance(dummy_input, dict):
                        model(**dummy_input)
                    else:
                        model(dummy_input)
                    
                    e.record()
                    torch.cuda.synchronize()
                    cur_model_infer_time = s.elapsed_time(e) / 1000.
                    infer_time_list += [cur_model_infer_time]

        else:
            with torch.no_grad():
                for _ in range(sample_num):
                    start = time.time()
                    
                    if isinstance(dummy_input, dict):
                        model(**dummy_input)
                    else:
                        model(dummy_input)
                    
                    cur_model_infer_time = time.time() - start
                    infer_time_list += [cur_model_infer_time]
                    
        avg_infer_time = sum(infer_time_list) / sample_num

        if return_detail:
            return avg_infer_time, infer_time_list
        return avg_infer_time