docstr: Shards the first MoE param and passes it to the underlying implementation. Since it's possible for an architecture
  to have both MoE and non-MoE layers, this can't be overloaded on the MLP1 transform. Furthermore, since both
  the MoE DSModule owns both MLP1 and MLP2, under certain sharding conditions it's not possible for the model implementation
  to infer from the shape whether to perform a different transformation based on MLP1 or MLP2. This (and the below)
  separations are intended to solve both these issues.
  
  Args:
      param (torch.Tensor): The parameter to transform. This should have shape (n_experts, out_neurons, in_neurons).
docstr_tok: ["Shards","the","first","MoE","param","and","passes","it","to","the","underlying","implementation",".","Since","it","'s","possible","for","an","architecture","to","have","both","MoE","and","non-MoE","layers",",","this","ca","n't","be","overloaded","on","the","MLP1","transform",".","Furthermore",",","since","both","the","MoE","DSModule","owns","both","MLP1","and","MLP2",",","under","certain","sharding","conditions","it","'s","not","possible","for","the","model","implementation","to","infer","from","the","shape","whether","to","perform","a","different","transformation","based","on","MLP1","or","MLP2",".","This","(","and","the","below",")","separations","are","intended","to","solve","both","these","issues",".","Args",":","param","(","torch.Tensor",")",":","The","parameter","to","transform",".","This","should","have","shape","(","n_experts",",","out_neurons",",","in_neurons",")","."]
code: def transform_moe_mlp_1_param(
      self, param: torch.Tensor
  ) -> InferenceParameter:
      """
      Shards the first MoE param and passes it to the underlying implementation. Since it's possible for an architecture
      to have both MoE and non-MoE layers, this can't be overloaded on the MLP1 transform. Furthermore, since both
      the MoE DSModule owns both MLP1 and MLP2, under certain sharding conditions it's not possible for the model implementation
      to infer from the shape whether to perform a different transformation based on MLP1 or MLP2. This (and the below)
      separations are intended to solve both these issues.
  
      Args:
          param (torch.Tensor): The parameter to transform. This should have shape (n_experts, out_neurons, in_neurons).
      """
      param = shard_mlp_1_param(
          param,
          self.tp_rank,
          self.tp_size,
          gated=self.gated_mlp,
          is_moe=True,
      )
  
      return self.moe.transform_moe_mlp_1_param(param)
  
imports: from abc import abstractmethod
  from typing import Optional
  import torch
  from deepspeed.accelerator import get_accelerator
  from ..config_v2 import RaggedInferenceEngineConfig
  from ..inference_utils import ActivationType, ceil_div, is_gated
  from ..model_implementations import *
  from ..model_implementations.sharding import *
  from ..modules.configs import DSEmbeddingsConfig, DSLinearConfig, DSMoEConfig, DSNormConfig, DSSelfAttentionConfig, DSUnembedConfig, NormTypeEnum, PositionalEmbeddingType, RotateHalfConfig
  from ..modules import heuristics
  from ..ragged import DSSequenceDescriptor, KVCacheConfig, RaggedBatchWrapper
  from .inference_model_base import DSInferenceModelBase, DSModelImplementationConfig, MPType
  from ..inference_parameter import InferenceParameter
label_desc: Use image node_id 6 for calling the DSMoETransformerModelBase obj's underlying member method code with example usage: obj.transform_moe_mlp_1_param(param) and returns: self