Update modeling_motif.py

modeling_motif.py

@@ -1,5 +1,5 @@
 import math
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, Callable, Dict
 
 import torch
 import torch.utils.checkpoint
@@ -35,144 +35,22 @@ logger = logging.get_logger(__name__)
 if is_flash_attn_2_available():
     from transformers.modeling_flash_attention_utils import _flash_attention_forward
 
-try:
-    moreh_ops = torch.ops.moreh
-    MorehRMSNorm = moreh_ops.T5LayerNorm
-    ScaledDotProductAttention = moreh_ops.scaled_dot_product_attention
-    MorehFlashAttention = moreh_ops.flash_attention
-    logger.warning_once("Using moreh ops")
-except AttributeError:
-    MorehRMSNorm = None
-    ScaledDotProductAttention = None
-    MorehFlashAttention = None
-    logger.warning_once("Failed to import moreh ops")
-
-
-
-# DEBUG = False
-# logger.info(f"DEBUG: {DEBUG} : will log timing")
-# def log_timing(obj):
-#     """Decorator to log timing of function or class execution"""
-#     if isinstance(obj, type):
-#         # If decorating a class
-#         class TimedClass(obj):
-#             def __getattribute__(self, name):
-#                 attr = super().__getattribute__(name)
-#                 if callable(attr) and not name.startswith('__'):
-#                     def timed_method(*args, **kwargs):
-#                         if not DEBUG:
-#                             return attr(*args, **kwargs)
-#                         if name != "forward":
-#                             return attr(*args, **kwargs)
-
-#                         start_time = time.time()
-#                         logger.info(f"Entering {obj.__name__}.{name}")
-#                         result = attr(*args, **kwargs)
-#                         end_time = time.time()
-#                         logger.info(f"Exiting {obj.__name__}.{name}, took {end_time - start_time:.4f} seconds")
-#                         return result
-#                     return timed_method
-#                 return attr
-#         return TimedClass
-#     else:
-#         # If decorating a function
-#         def wrapper(*args, **kwargs):
-#             if not DEBUG:
-#                 return obj(*args, **kwargs)
-
-#             start_time = time.time()
-#             logger.info(f"Entering {obj.__name__}")
-#             result = obj(*args, **kwargs)
-#             end_time = time.time()
-#             logger.info(f"Exiting {obj.__name__}, took {end_time - start_time:.4f} seconds")
-#             return result
-#         return wrapper
-
-
+MorehRMSNorm = None
+ScaledDotProductAttention = None
+MorehFlashAttention = None
 
 #_CHECKPOINT_FOR_DOC = "moreh/Motif-102B"
 _CONFIG_FOR_DOC = "MotifConfig"
 
-#from .moreh_moe import MorehMoeMLP, MorehMoeFusedMLP
-
-import torch
 from transformers.activations import ACT2CLS as _ACT2CLS
 from transformers.activations import ClassInstantier
-moreh_ops = torch.ops.moreh
-
-from typing import Callable, Dict, List, Tuple
-
-import torch
-
-
-# @log_timing
-def multi_head_forward_backward(shared_activation: torch.Tensor,
-                                head_fns: List[Callable[[torch.Tensor], Dict[str, torch.Tensor]]],
-                                return_keys=("loss", ),
-                                return_only_first_head=True) -> Tuple[torch.Tensor, ...]:
-    """
-    The forward-backward pattern introduced in the paper https://arxiv.org/abs/2404.19737
-    to reduce memory overhead due to activations from multiple heads.
-
-    Args:
-    - shared_activation: the shared activation across all heads
-    - head_fns: the head-wise forward computations that start from `shared_activation`.
-        it should output a dictionary of tensors with keys matching `return_keys`
-    - return_keys: the keys to return in order
-    - return_only_first_head: whether to return only the values from the first head
-
-    Returns:
-    - a tuple of return tensors
-
-    Side effect:
-    - (only when `torch.is_grad_enabled()`)
-        the gradients accumulated as if `sum(head_fn(shared_activation)["loss"] for head_fn in head_fns).backward()` had been called
-    """
-    if not return_only_first_head:
-        raise NotImplementedError
-
-    return_key_set = set(return_keys)
-    if "loss" not in return_key_set:
-        raise Exception("'loss' is a required return key.")
-
-    detached_shared_activation = shared_activation.detach()
-    detached_shared_activation.requires_grad = True
-    return_values = {key: None for key in return_keys}
-    for head_idx, head_fn in enumerate(head_fns):
-        if head_idx > 0 and not torch.is_grad_enabled():
-            continue
-
-        # forward pass for the head
-        headwise_outputs = head_fn(detached_shared_activation)
-        if set(headwise_outputs.keys()) != return_key_set:
-            raise Exception(f"Headwise output keys {headwise_outputs.keys()} do not match return keys {return_keys}.")
-
-        # backward pass for the head
-        # effect 1: the parameters of the head
-        # effect 2: gradient accumulated in `detached_shared_activation.grad`
-        if torch.is_grad_enabled():
-            headwise_loss = headwise_outputs["loss"]
-            headwise_loss.backward(
-            )  # NOTE: You do not need to retain graph since no graph is shared across backward passes
-
-        if head_idx == 0:
-            for key in return_keys:
-                return_values[key] = headwise_outputs[key]
-
-    assert all(value is not None for value in return_values.values())
-
-    # backward pass for the shared part
-    if torch.is_grad_enabled():
-        shared_activation.backward(detached_shared_activation.grad)
-
-    return tuple(return_values[key] for key in return_keys)
 
 
 class PolyNorm(torch.nn.Module):
     """
     A trainable activation function introduced in https://arxiv.org/html/2411.03884v1.
     The code is copied from https://github.com/BryceZhuo/PolyCom?tab=readme-ov-file/README.md,
-    with the change `* torch.rsqrt` => `/ torch.sqrt` for potential MAF incompatibility.
+    with the change `* torch.rsqrt` => `/ torch.sqrt`
     """
 
     def __init__(self, eps=1e-6):
@@ -189,31 +67,11 @@ class PolyNorm(torch.nn.Module):
             x ** 2) + self.weight[2] * self._norm(x) + self.bias
 
 
-class PolyNorm_Test(torch.nn.Module):
-    """
-    A trainable activation function introduced in https://arxiv.org/html/2411.03884v1.
-    The code is copied from https://github.com/BryceZhuo/PolyCom?tab=readme-ov-file/README.md,
-    with the change `* torch.rsqrt` => `/ torch.sqrt` for potential MAF incompatibility.
-    """
-
-    def __init__(self, eps=1e-6):
-        super(PolyNorm_Test, self).__init__()
-        self.weight = torch.nn.Parameter(torch.ones(3) / 3)
-        self.bias = torch.nn.Parameter(torch.zeros(1))
-        self.eps = eps
-
-    def forward(self, x):
-
-        #return torch.nn.SiLU(x)
-        return moreh_ops.poly_norm(x, self.weight, self.bias)
-
-
-CUSTOM_ACT2CLS = {"poly_norm": PolyNorm_Test, "poly_norm_test": PolyNorm_Test}
+CUSTOM_ACT2CLS = {"poly_norm": PolyNorm}
 ACT2CLS = {**_ACT2CLS, **CUSTOM_ACT2CLS}
 ACT2FN = ClassInstantier(ACT2CLS)
 
 
-
 class MotifRMSNorm(nn.Module):
 
     def __init__(self, hidden_size, eps=1e-6):
@@ -226,8 +84,7 @@ class MotifRMSNorm(nn.Module):
 
     def forward(self, hidden_states):
         input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         return self.weight * hidden_states.to(input_dtype)
 
@@ -235,7 +92,7 @@ class MotifRMSNorm(nn.Module):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
-ALL_LAYERNORM_LAYERS.append(MotifRMSNorm if MorehRMSNorm is None else MorehRMSNorm)
+ALL_LAYERNORM_LAYERS.append(MotifRMSNorm)
 
 
 class MotifRotaryEmbeddingWithCache(nn.Module):
@@ -293,7 +150,6 @@ class MotifRotaryEmbeddingWithCache(nn.Module):
         )
 
 
-# @log_timing
 class MotifRotaryEmbedding(nn.Module):
 
     def __init__(
@@ -307,7 +163,6 @@ class MotifRotaryEmbedding(nn.Module):
         config: Optional[MotifConfig] = None,
     ):
         super().__init__()
-        # TODO (joao): remove the `if` below, only used for BC
         self.rope_kwargs = {}
         if config is None:
             logger.warning_once(
@@ -352,7 +207,7 @@ class MotifRotaryEmbedding(nn.Module):
                                                                  device,
                                                                  seq_len=seq_len,
                                                                  **self.rope_kwargs)
-            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilation
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
             self.max_seq_len_cached = seq_len
 
         if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
@@ -401,7 +256,6 @@ def rotate_half(x):
     return rotated_tensor
 
 
-# @log_timing
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1, fused_rope=True):
     """
     Applies rotary position embeddings to the input tensors.
@@ -427,12 +281,7 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1, fus
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
     '''
-    #cos = cos[position_ids]
-    #sin = sin[position_ids]
-
-    #cos = cos[position_ids].unsqueeze(unsqueeze_dim)  # [bs, 1, seq_len, dim]
-    #sin = sin[position_ids].unsqueeze(unsqueeze_dim)  # [bs, 1, seq_len, dim]
-
+    
     q = q.transpose(1, 2)
     k = k.transpose(1, 2)
 
@@ -450,9 +299,8 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1, fus
     return q_embed, k_embed
 
 
-# @log_timing
 class MotifMLP(nn.Module):
-
+    
     def __init__(self, config):
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -462,394 +310,11 @@ class MotifMLP(nn.Module):
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
         self.act_fn = ACT2FN[config.hidden_act]
 
-        if config.wesar_weights:
-            self.gate_up_proj_alpha = nn.Parameter(torch.tensor(1) *config.gate_up_proj_alpha)
-            self.down_proj_alpha = nn.Parameter(torch.tensor(1) * config.down_proj_alpha)
-        else:
-            self.gate_up_proj_alpha=1
-            self.down_proj_alpha=1
-        if config.muP:
-            self.down_proj.__do_scale_tager__ = True
-            self.gate_proj.__do_scale_tager_mu_dim_model__  = True
-            self.up_proj.__do_scale_tager_mu_dim_model__ = True
-            self.down_proj.__do_scale_tager_mu_ffn__ = True
-
-
     def forward(self, hidden_state):
-        hidden_state = hidden_state*self.gate_up_proj_alpha
-        #hidden_state = self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))*
-        return self.down_proj_alpha*self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
-
-
-class MorehMoeFusedMLP(nn.Module):
-    def __init__(self,
-                 ffn_dim,
-                 hidden_dim,
-                 hidden_act_moe,
-                 num_experts,
-                 num_groups=1,
-                 device=None,
-                 continual_training=False):
-        super().__init__()
-        self.ffn_dim = ffn_dim
-        self.hidden_dim = hidden_dim
-        self.hidden_act_moe = hidden_act_moe
-
-        self.num_experts = num_experts
-        self.num_groups = num_groups
-
-        assert self.num_experts % self.num_groups == 0
-        self.num_experts_per_group = self.num_experts // self.num_groups
-
-        ## bsz, seq, group size, 2*ffn_size
-
-        moreh_ops = torch.ops.moreh
-        self.w13 = nn.ModuleList([
-            moreh_ops.MoeFanInLinear(self.hidden_dim,
-                                     self.ffn_dim * 2,
-                                     bias=False,
-                                     num_experts=self.num_experts_per_group,
-                                     device=device)
-            for _ in range(self.num_groups)
-        ])
-
-        self.w2 = nn.ModuleList([
-            moreh_ops.MoeFanOutLinear(self.ffn_dim,
-                                      self.hidden_dim,
-                                      bias=False,
-                                      num_experts=self.num_experts_per_group,
-                                      device=device)
-            for _ in range(self.num_groups)
-        ])
-
-        ## use silu?
-        self.act_fn = ACT2FN[self.hidden_act_moe]
-
-        if continual_training:
-            logger.info('two optipons 1. zero init all weights, 2. add scaling param to moe output.')
-            self._zero_init()
-
-    def _zero_init(self):
-        for module in self.w2:
-            for n,param in module.named_parameters():
-                logger.info(f'{n} {param.shape}')
-                param.data.zero_()
-
-
-    def forward(self, hidden_states, selected_experts, routing_weights):
-        w13_final_output = None
-        for group_idx in range(self.num_groups):
-            w13_output_in_group = self._get_w13_output(hidden_states,
-                                                       selected_experts,
-                                                       group_idx)
-            if w13_final_output is None:
-                w13_final_output = w13_output_in_group
-            else:
-                w13_final_output += w13_output_in_group
-
-        current_hidden_states = self.act_fn(
-            w13_final_output[:, :, :, :self.ffn_dim]
-        ) * w13_final_output[:, :, :, self.ffn_dim:]
-
-        final_hidden_states = None
-        for group_idx in range(self.num_groups):
-            w2_output_in_group = self._get_w2_output(current_hidden_states,
-                                                     selected_experts,
-                                                     routing_weights, group_idx)
-            if final_hidden_states is None:
-                final_hidden_states = w2_output_in_group
-            else:
-                final_hidden_states += w2_output_in_group
-        return final_hidden_states
-
-    def _get_w13_output(self, hidden_states, selected_experts, group_idx):
-        selected_experts_in_group = selected_experts - (
-            group_idx * self.num_experts_per_group)
-
-        w13_output = self.w13[group_idx](hidden_states,
-                                         selected_experts_in_group)
-        return w13_output
-
-    def _get_w2_output(self, hidden_states, selected_experts, routing_weights,
-                       group_idx):
-        selected_experts_in_group = selected_experts - (
-            group_idx * self.num_experts_per_group)
-        output = self.w2[group_idx](hidden_states, selected_experts_in_group,
-                                    routing_weights)
-        return output
-
-
-class MoEGate(nn.Module):
-
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.top_k = config.num_experts_per_tok
-        self.n_routed_experts = config.n_routed_experts
-        self.routed_scaling_factor = config.routed_scaling_factor
-        self.scoring_func = config.scoring_func
-        self.seq_aux = config.seq_aux
-        self.topk_method = config.topk_method
-        self.n_group = config.n_group
-        self.topk_group = config.topk_group
-
-        # topk selection algorithm
-        self.norm_topk_prob = config.norm_topk_prob
-        self.gating_dim = config.hidden_size
-        self.weight = nn.Parameter(
-            torch.empty((self.n_routed_experts, self.gating_dim)))
-        if self.topk_method == "noaux_tc":
-            self.e_score_correction_bias = nn.Parameter(
-                torch.empty((self.n_routed_experts)))
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        import torch.nn.init as init
-
-        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
-
-    def forward(self, hidden_states):
-        bsz, seq_len, h = hidden_states.shape
-        ### compute gating score
-        hidden_states = hidden_states.view(-1, h)
-        logits = F.linear(hidden_states.type(torch.float32),
-                          self.weight.type(torch.float32), None)
-        if self.scoring_func == "sigmoid":
-            scores = logits.sigmoid()
-        else:
-            raise NotImplementedError(
-                f"insupportable scoring function for MoE gating: {self.scoring_func}"
-            )
-
-        ### select top-k experts
-        if self.topk_method == "greedy":
-            topk_weight, topk_idx = torch.topk(scores,
-                                               k=self.top_k,
-                                               dim=-1,
-                                               sorted=False)
-        elif self.topk_method == "group_limited_greedy":
-            group_scores = (scores.view(bsz * seq_len, self.n_group,
-                                        -1).max(dim=-1).values)  # [n, n_group]
-            group_idx = torch.topk(group_scores,
-                                   k=self.topk_group,
-                                   dim=-1,
-                                   sorted=False)[1]  # [n, top_k_group]
-            group_mask = torch.zeros_like(group_scores)  # [n, n_group]
-            group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
-            score_mask = (group_mask.unsqueeze(-1).expand(
-                bsz * seq_len, self.n_group,
-                self.n_routed_experts // self.n_group).reshape(
-                    bsz * seq_len, -1))  # [n, e]
-            tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0)  # [n, e]
-            topk_weight, topk_idx = torch.topk(tmp_scores,
-                                               k=self.top_k,
-                                               dim=-1,
-                                               sorted=False)
-        elif self.topk_method == "noaux_tc":
-            ###  will be used. ###
-            scores_for_choice = scores.view(
-                bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0)
-            group_scores = (scores_for_choice.view(
-                bsz * seq_len, self.n_group,
-                -1).topk(2, dim=-1)[0].sum(dim=-1))  # [n, n_group]
-            group_idx = torch.topk(group_scores,
-                                   k=self.topk_group,
-                                   dim=-1,
-                                   sorted=False)[1]  # [n, top_k_group]
-            group_mask = torch.zeros_like(group_scores)  # [n, n_group]
-            group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
-            score_mask = (group_mask.unsqueeze(-1).expand(
-                bsz * seq_len, self.n_group,
-                self.n_routed_experts // self.n_group).reshape(
-                    bsz * seq_len, -1))  # [n, e]
-            tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(),
-                                                       0.0)  # [n, e]
-            _, topk_idx = torch.topk(tmp_scores,
-                                     k=self.top_k,
-                                     dim=-1,
-                                     sorted=False)
-            topk_weight = scores.gather(1, topk_idx)
-        else:
-            raise NotImplementedError(
-                f"insupportable TopK function for MoE gating: {self.topk_method}"
-            )
-
-        ### norm gate to sum 1
-        if self.top_k > 1 and self.norm_topk_prob:
-            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
-            topk_weight = topk_weight / denominator
-        topk_weight = topk_weight * self.routed_scaling_factor  # must multiply the scaling factor
-
-        return topk_idx, topk_weight
-
-
-class MotifMoE(nn.Module):
-    """
-    A mixed expert module containing shared experts.
-    """
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.num_experts_per_tok = config.num_experts_per_tok
-        self.use_moreh_moe = config.use_moreh_moe
-        self.use_fused_mlp = config.use_fused_mlp
-
-        if hasattr(config, "ep_size") and config.ep_size > 1:
-            assert config.ep_size == dist.get_world_size()
-            assert not config.use_moreh_moe
-            self.ep_size = config.ep_size
-            self.experts_per_rank = config.n_routed_experts // config.ep_size
-            self.ep_rank = dist.get_rank()
-            self.experts = nn.ModuleList([
-                (DeepseekV3MLP(config,
-                               intermediate_size=config.moe_intermediate_size)
-                 if i >= self.ep_rank * self.experts_per_rank and i <
-                 (self.ep_rank + 1) * self.experts_per_rank else None)
-                for i in range(config.n_routed_experts)
-            ])
-        else:
-            self.ep_size = 1
-            self.experts_per_rank = config.n_routed_experts
-            self.ep_rank = 0
-            if self.use_moreh_moe:
-                if not self.use_fused_mlp:
-                    self.experts = MorehMoeMLP(
-                        ffn_dim=config.moe_intermediate_size,
-                        hidden_dim=config.hidden_size,
-                        hidden_act_moe=config.hidden_act_moe,
-                        num_experts=config.n_routed_experts,
-                        device=None)
-                else:
-                    ## group expert.
-                    self.experts = MorehMoeFusedMLP(
-                        ffn_dim=config.moe_intermediate_size,
-                        hidden_dim=config.hidden_size,
-                        hidden_act_moe=config.hidden_act_moe,
-                        num_experts=config.n_routed_experts,
-                        num_groups=config.n_group,
-                        device=None,
-                        continual_training=config.continual_training,
-                        )
-            else:
-                self.experts = nn.ModuleList([
-                    DeepseekV3MLP(
-                        config, intermediate_size=config.moe_intermediate_size)
-                    for i in range(config.n_routed_experts)
-                ])
-
-        self.gate = MoEGate(config)
-
-    def forward(self, hidden_states):
-        identity = hidden_states
-        orig_shape = hidden_states.shape
-        topk_idx, topk_weight = self.gate(hidden_states)
-        if self.use_moreh_moe:
-            y = self.experts(hidden_states, topk_idx.view(*orig_shape[:-1], -1),
-                             topk_weight.view(*orig_shape[:-1], -1))
-            y = y.type(hidden_states.dtype)
-        else:
-            hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-            flat_topk_idx = topk_idx.view(-1)
-            if self.training:
-                hidden_states = hidden_states.repeat_interleave(
-                    self.num_experts_per_tok, dim=0)
-                y = torch.empty_like(hidden_states)
-                for i, expert in enumerate(self.experts):
-                    y[flat_topk_idx == i] = expert(
-                        hidden_states[flat_topk_idx == i])
-                y = (y.view(*topk_weight.shape, -1) *
-                     topk_weight.unsqueeze(-1)).sum(dim=1)
-                y = y.type(hidden_states.dtype)
-                y = y.view(*orig_shape)
-                # y = AddAuxiliaryLoss.apply(y, aux_loss)
-            else:
-                y = self.moe_infer(hidden_states, topk_idx,
-                                   topk_weight).view(*orig_shape)
-        return y, identity
-
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens = x[idxs // topk_ids.shape[1]]
-        sorted_tokens_shape = sorted_tokens.shape
-        if self.ep_size > 1:
-            tokens_per_ep_rank = tokens_per_expert.view(self.ep_size,
-                                                        -1).sum(dim=1)
-            tokens_per_expert_group = tokens_per_expert.new_empty(
-                tokens_per_expert.shape[0])
-            dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert)
-            output_splits = (tokens_per_expert_group.view(
-                self.ep_size, -1).sum(1).cpu().numpy().tolist())
-            gathered_tokens = sorted_tokens.new_empty(
-                tokens_per_expert_group.sum(dim=0).cpu().item(),
-                sorted_tokens.shape[1])
-            input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist()
-            dist.all_to_all(
-                list(gathered_tokens.split(output_splits)),
-                list(sorted_tokens.split(input_split_sizes)),
-            )
-            tokens_per_expert_post_gather = tokens_per_expert_group.view(
-                self.ep_size, self.experts_per_rank).sum(dim=0)
-            gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],),
-                                    dtype=np.int32)
-            s = 0
-            for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
-                gatherd_idxs[s:s + k] = i % self.experts_per_rank
-                s += k
-            gatherd_idxs = gatherd_idxs.argsort()
-            sorted_tokens = gathered_tokens[gatherd_idxs]
-            tokens_per_expert = tokens_per_expert_post_gather
-        tokens_per_expert = tokens_per_expert.cpu().numpy()
-
-        outputs = []
-        start_idx = 0
-        for i, num_tokens in enumerate(tokens_per_expert):
-            end_idx = start_idx + num_tokens
-            if num_tokens == 0:
-                continue
-            expert = self.experts[i + self.ep_rank * self.experts_per_rank]
-            tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
-            expert_out = expert(tokens_for_this_expert)
-            outputs.append(expert_out)
-            start_idx = end_idx
-
-        outs = torch.cat(outputs,
-                         dim=0) if len(outputs) else sorted_tokens.new_empty(0)
-        if self.ep_size > 1:
-            new_x = torch.empty_like(outs)
-            new_x[gatherd_idxs] = outs
-            gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
-            dist.all_to_all(
-                list(gathered_tokens.split(input_split_sizes)),
-                list(new_x.split(output_splits)),
-            )
-            outs = gathered_tokens
-
-        new_x = torch.empty_like(outs)
-        new_x[idxs] = outs
-        final_out = (new_x.view(
-            *topk_ids.shape, -1).type(topk_weight.dtype).mul_(
-                topk_weight.unsqueeze(dim=-1)).sum(dim=1).type(new_x.dtype))
-        return final_out
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
 
 
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-
-
-    """
-    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
-    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
-
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
-    if n_rep == 1:
-        return hidden_states
-    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-    """
-
     return torch.repeat_interleave(hidden_states, dim=1, repeats=n_rep)
 
 
@@ -1384,7 +849,6 @@ MOTIF_ATTENTION_CLASSES = {
 }
 
 
-# @log_timing
 class MotifDecoderLayer(nn.Module):
 
     def __init__(self, config: MotifConfig, moe_layer: bool, layer_idx: int):
@@ -1655,7 +1119,6 @@ MOTIF_INPUTS_DOCSTRING = r"""
 """
 
 
-# @log_timing
 @add_start_docstrings(
     "The bare Motif Model outputting raw hidden-states without any specific head on top.",
     MOTIF_START_DOCSTRING,
@@ -1675,17 +1138,13 @@ class MotifModel(MotifPreTrainedModel):
         self.multi_token_heads = config.multi_token_heads
 
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
-        # NOTE: For multi-token models, the last decoder layers (one for each token index)
-        # are implemented as a part of `MotifModelForCausalLM` to enable a custom forward-backward procedure.
 
         num_hidden_layers = config.num_hidden_layers if self.multi_token_heads is None else config.num_hidden_layers - 1
-        if config.moe:
-            moe_layer = [True for i in range(num_hidden_layers)]
-        else:
-            moe_layer = [False for i in range(num_hidden_layers)]
+
         logger.info(f'current_moe layer { moe_layer }')
-        self.layers = nn.ModuleList([MotifDecoderLayer(config = config, moe_layer= moe_layer[layer_idx],
-                                    layer_idx=layer_idx) for layer_idx in range(num_hidden_layers)])
+        self.layers = nn.ModuleList([
+            MotifDecoderLayer(config = config, layer_idx=layer_idx) for layer_idx in range(num_hidden_layers)
+        ])
         self._attn_implementation = config._attn_implementation
         RMSNorm = MorehRMSNorm if MorehRMSNorm is not None else MotifRMSNorm
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -1701,36 +1160,6 @@ class MotifModel(MotifPreTrainedModel):
         self.gradient_checkpointing = False
         self.post_init()
 
-        self.use_pipeline = config.use_pipeline
-        if self.use_pipeline:
-            logger.info('use reinforced pp..')
-            if config.num_stages==2:
-                ### moe version
-                if config.decontam_attn:
-                    self.split_layers = [15]
-                else:
-                    if num_hidden_layers == 32:
-                        self.split_layers = [15] # 14: 15,17 # 13: 14:18
-                    else:
-                        self.split_layers = [6]
-            elif config.num_stages==3:
-                self.split_layers = [9,20] ## 11, 11, 10
-            elif config.num_stages==4:
-                self.split_layers = [7,15,23] #7,9,9,7
-            elif config.num_stages==16:
-                self.split_layers = [1,3,5,7,9,11,13,15,17,19,21,23,25,27,29]
-            logger.info(f' check the split layers (moe): {self.split_layers}')
-
-        self.scale_emb = 1
-
-        # Reparameterization <|_1_|>
-        if config.wesar_weights :
-            logger.info(f'config.wesar_weights {config.wesar_weights}')
-            self.norm_alpha = nn.Parameter(torch.tensor(1).float())
-            self.scale_emb = 10
-        else:
-            self.norm_alpha = 1
-
     def get_input_embeddings(self):
         return self.embed_tokens
 
@@ -1769,7 +1198,6 @@ class MotifModel(MotifPreTrainedModel):
                     "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...")
                 use_cache = False
 
-        # kept for BC (non `Cache` `past_key_values` inputs)
         return_legacy_cache = False
         if use_cache and not isinstance(past_key_values, Cache):
             return_legacy_cache = True
@@ -1783,7 +1211,7 @@ class MotifModel(MotifPreTrainedModel):
                     "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)")
 
         if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids) * self.scale_emb
+            inputs_embeds = self.embed_tokens(input_ids)
 
         if cache_position is None:
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
@@ -1837,18 +1265,13 @@ class MotifModel(MotifPreTrainedModel):
 
             hidden_states = layer_outputs[0]
 
-
-            if self.use_pipeline and idx in self.split_layers:
-                hidden_states = torch.moreh.pipeline_assign(hidden_states)
-
             if use_cache:
                 next_decoder_cache = layer_outputs[2 if output_attentions else 1]
 
             if output_attentions:
                 all_self_attns += (layer_outputs[1], )
 
-        # <|_2_|>
-        hidden_states = self.norm(hidden_states)* self.norm_alpha
+        hidden_states = self.norm(hidden_states)
 
         # add hidden states from the last decoder layer
         if output_hidden_states:
@@ -1881,8 +1304,6 @@ class MotifModel(MotifPreTrainedModel):
         output_attentions: bool,
     ):
         if self.config._attn_implementation == "flash_attention_2":
-            if MorehFlashAttention is not None:
-                return attention_mask
             if attention_mask is not None and 0.0 in attention_mask:
                 return attention_mask
             return None
@@ -2003,7 +1424,6 @@ class MotifModel(MotifPreTrainedModel):
         return causal_mask
 
 
-# @log_timing
 class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
 
@@ -2013,33 +1433,19 @@ class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
         self.vocab_size = config.vocab_size
         self.multi_token_heads = config.multi_token_heads
 
-        if self.multi_token_heads is None:
-            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         else:
             self.tokenwise_last_layers = nn.ModuleList(
                 [MotifDecoderLayer(config, config.num_hidden_layers - 1) for _ in range(self.multi_token_heads)])
             self.tokenwise_lm_heads = nn.ModuleList(
                 [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(self.multi_token_heads)])
-        self.should_skip_separate_backward_pass = self.multi_token_heads is not None
 
         # Initialize weights and apply final processing
         self.post_init()
-
-        # <|_3_|>
-        if config.muP:
-            self.lm_head.__do_scale_tager_mu_dim_base_model__=True
-
-        # <|_4_|>
-        self.lm_head_alpha = 1
-        if config.wesar_weights:
-            self.lm_head_alpha = nn.Parameter(torch.tensor(1).float())
-
+        
         if getattr(config, "tie_word_embeddings", True):
             logger.info('tie embeddings')
             self.tie_weights()
-        else:
-            # <|_5_|>
-            self.lm_head.__do_scale_tager_mu_dim_base_model__ = False
 
     def get_input_embeddings(self):
         return self.model.embed_tokens
@@ -2059,101 +1465,7 @@ class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
     def get_decoder(self):
         return self.model
 
-    def multi_token_forward_backward(self,
-                                     hidden_states: torch.FloatTensor,
-                                     outputs: MotifModelOutputWithPast,
-                                     labels: torch.LongTensor,
-                                     position_ids: Optional[torch.LongTensor],
-                                     output_attentions: Optional[bool],
-                                     use_cache: Optional[bool],
-                                     cache_position: Optional[torch.LongTensor],
-                                     return_dict: Optional[bool],
-                                     num_logits_to_keep: int = 0) -> CausalLMOutputWithPast:
-        """
-        This implements the main forward-backward procedure for multi-token model training proposed in
-        the paper https://arxiv.org/abs/2404.19737.
-        Essentially,
-        - The multi-token model tries to predict n (instead of 1) tokens at a time.
-        - Applying this only during training and using first-token prediction during inference is still helpful.
-        - The change in architecture: when using n-token prediction, each token index (between 1 and n) has its own
-            (1) last attention layer and (2) lm head.
-        - The change in loss: sum of cross-entropy losses corresponding to each token index.
-        - Custom forward-backward procedure for memory efficiency: refer to the implementation of `multi_head_forward_backward`.
-        """
-        if not return_dict:
-            raise NotImplementedError("return_dict must be True for multi-token training")
-
-        past_key_values = outputs.past_key_values
-        causal_mask = outputs.causal_mask
-        position_embeddings = outputs.position_embeddings
-
-        if labels is not None:
-            labels = labels.to(hidden_states.device)
-
-        def _tokenwise_forward(hidden_states: torch.Tensor, token_idx):
-            ## Model forward
-            layer = self.tokenwise_last_layers[token_idx]
-            lm_head = self.tokenwise_lm_heads[token_idx]
-
-            layer_outputs = layer(
-                hidden_states,
-                attention_mask=causal_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,  # TODO: update past_key_values?
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cache_position=cache_position,
-                position_embeddings=position_embeddings,
-            )
-            last_hidden_states = layer_outputs[0]
-            if num_logits_to_keep > 0:
-                assert labels is None
-                last_hidden_states = last_hidden_states[:, -num_logits_to_keep:, :]
-            tokenwise_logits = lm_head(last_hidden_states)
-
-            if labels is None:
-                return {
-                    "loss": None,
-                    "logits": tokenwise_logits,
-                }
-
-            ## Compute loss
-            shift_n = token_idx + 1
-            shift_logits = tokenwise_logits[..., :-shift_n, :].contiguous()
-            shift_labels = labels[..., shift_n:].contiguous()
-
-            loss_fct = CrossEntropyLoss()
-            shift_logits = shift_logits.view(-1, self.config.vocab_size)
-            shift_labels = shift_labels.view(-1)
-
-            tokenwise_loss = loss_fct(shift_logits, shift_labels)
-
-            return {
-                "loss": tokenwise_loss,
-                "logits": tokenwise_logits,
-            }
-
-        head_fns = [
-            lambda hidden_states, token_idx=token_idx: _tokenwise_forward(hidden_states, token_idx)
-            for token_idx in range(self.multi_token_heads)
-        ]
-        loss, logits = multi_head_forward_backward(hidden_states,
-                                                   head_fns,
-                                                   return_keys=("loss", "logits"),
-                                                   return_only_first_head=True)
-
-        if not return_dict:
-            output = (logits, ) + outputs[1:]
-            return (loss, ) + output
-
-        return CausalLMOutputWithPast(
-            loss=loss,
-            logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
+    
     @add_start_docstrings_to_model_forward(MOTIF_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(
@@ -2191,8 +1503,8 @@ class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
         ```python
         >>> from transformers import AutoTokenizer, MotifForCausalLM
 
-        >>> model = MotifForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
-        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+        >>> model = MotifForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS, trust_remote_code = True)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER, trust_remote_code = True)
 
         >>> prompt = "Hey, are you conscious? Can you talk to me?"
         >>> inputs = tokenizer(prompt, return_tensors="pt")
@@ -2209,8 +1521,6 @@ class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-        outputs_include_causal_mask = self.multi_token_heads is not None
-        outputs_include_position_embeddings = self.multi_token_heads is not None
         outputs: MotifModelOutputWithPast = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -2222,31 +1532,16 @@ class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
             cache_position=cache_position,
-            outputs_include_causal_mask=outputs_include_causal_mask,
-            outputs_include_position_embeddings=outputs_include_position_embeddings,
         )
 
         hidden_states = outputs[0]
 
-        if self.multi_token_heads is not None:
-            return self.multi_token_forward_backward(hidden_states,
-                                                     outputs,
-                                                     labels,
-                                                     position_ids,
-                                                     output_attentions,
-                                                     use_cache,
-                                                     cache_position,
-                                                     return_dict,
-                                                     num_logits_to_keep=num_logits_to_keep)
-
         # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
-        hidden_states = hidden_states * self.lm_head_alpha
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
         logits = logits.float()
 
         loss = None
         if labels is not None:
-            logits = logits
             # Shift so that tokens < n predict n
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()