马斯克xAI开源3140亿参数模型grok，将近2倍GPT3模型参数

马斯克说若OpenAI更名为closeAI，其针对OpenAI的起诉讲会被撤销。很是显然马斯克针对OpenAI的创业的初衷有一定的质疑。但是无论如何OpenAI也无法更名为为close AI，其公司下的ChatGPT以及文生视频大模型Sora也不会开源。但是随着马斯克旗下公司xAI大模型的成功，其马斯克也兑现了自己开源的承诺，把3140亿参数的grok模型进行了开源。

Grok 是一款模仿《银河系漫游指南》的人工智能，因此几乎可以回答任何问题，其grok的动力引擎便是grok-1模型，此模型拥有3140亿个训练参数。而GPT-3模型的参数是1750亿参数，差不多1.8倍左右。Grok-1 是一个由 xAI 从头开始​训练的 3140 亿参数混合专家模型。Grok-1 预训练阶段在 2023 年 10 月结束，这意味着该模型没有针对任何特定应用（例如对话）进行微调。grok在 Apache 2.0 许可下发布权重和架构。

基础模型基于大量文本数据进行训练，没有针对任何具体任务进行微调；3140 亿参数的 MoE 模型，在给定 token 上的激活权重为 25%；2023 年 10 月，xAI 使用 JAX 库和 Rust 语言组成的自定义训练堆栈从头开始训练

模型参数已经超越了ChatGPT，其性能是否也同样超越了其他开源或者流行的LLM大语言模型。grok使用了一些数据集进行了训练并与其他模型进行了对比，在一些特定的学科以及问题上面，其grok的性能超越了LIama2 70B，ChatGPT等模型。

从开源代码可以看出，其模型是基于transformer的混合专家MoE模型。模型输入此表128*1024，输入序列长度为8192，embedding尺寸为48*128，多头注意力机制头数为48，一共有64层的解码器层，使用了8个专家模型，在选择专家模型时，使用2个专家进行结果的输出。

def main(): grok_1_model = LanguageModelConfig( vocab_size=128 * 1024, pad_token=0, eos_token=2, sequence_len=8192, embedding_init_scale=1.0, output_multiplier_scale=0.5773502691896257, embedding_multiplier_scale=78.38367176906169, model=TransformerConfig( emb_size=48 * 128, widening_factor=8, key_size=128, num_q_heads=48, num_kv_heads=8, num_layers=64, attn_output_multiplier=0.08838834764831845, shard_activations=True, # MoE. num_experts=8, num_selected_experts=2, # Activation sharding. data_axis="data", model_axis="model", ),

虽然模型已经开源，且预训练模型权重已经提供下载链接，但是3140亿参数的大模型不是一般电脑可以运行起来的。这也是为什么大语言模型一旦训练完成，就没有办法基于下游任务进行微调，其微调成本确实是巨大的，没有几个电脑可以运行此大模型。当然若运行此大模型到底需要多少的配置，初步估计需要8*80GB的H100。

而在词嵌入embedding处理与标准的transformer的直接嵌入有许些差别，其grok采用了一种称之为旋转嵌入的方式来进行词嵌入操作。

class RotaryEmbedding(hk.Module): """Applies rotary embeddings (RoPE) to the input sequence tensor, as described in https://arxiv.org/abs/2104.09864. """ def __init__( self, dim: int, name: Optional[str] = None, base_exponent: int = 10000, ): super().__init__(name) self.dim = dim self.base_exponent = base_exponent assert self.dim % 2 == 0 def __call__( self, x: jax.Array, seq_dim: int, offset: jax.Array, const_position: Optional[int] = None, t: Optional[jax.Array] = None, ) -> jax.Array: fprop_dtype = x.dtype # Compute the per-dimension frequencies exponents = jnp.arange(0, self.dim, 2, dtype=jnp.float32) inv_freq = jnp.asarray( 1.0 / (self.base_exponent ** (exponents / self.dim)), dtype=jnp.float32 ) if jnp.shape(offset) == (): # Offset can be a scalar or one offset per batch element. offset = jnp.expand_dims(offset, 0) # Compute the per element phase (to pass into sin and cos) if const_position: t = const_position * jnp.ones( ( 1, x.shape[seq_dim], ), dtype=jnp.float32, ) elif t is None: t = jnp.arange(x.shape[seq_dim], dtype=jnp.float32) + jnp.expand_dims(offset, -1) phase = jnp.einsum("bi,j->bij", t, inv_freq) phase = jnp.tile(phase, reps=(1, 2))[:, :, None, :] x = x * jnp.cos(phase) + rotate_half(x) * jnp.sin(phase) x = x.astype(fprop_dtype) return x

当然模型整体框架基于transformer模型的解码器部分，其混合了专家MoE模型。混合专家模型（Mixture of Experts，MoE）是一种机器学习模型，它结合了多个专家模型的输出以进行预测。这种模型的基本思想是，不同的专家模型擅长处理不同类型的输入数据，因此通过将这些专家模型的输出进行混合，可以提高模型的整体性能。

class MoELayer(hk.Module): def __init__( self, num_experts: int, layer_fn: Callable, router: Router, mesh: Any = None, shard_activations: bool = False, data_axis: Union[str, Tuple[str, ...]] = "data", model_axis: Union[str, Tuple[str, ...]] = "model", name: Optional[str] = "moe", ): super().__init__(name) self.num_experts = num_experts self.layer_fn = layer_fn self.router = router self.mesh = mesh self.shard_activations = shard_activations self.data_axis = data_axis self.model_axis = model_axis

混合专家模型通常由两部分组成：专家网络和门控网络。专家网络由多个专家模型组成，每个专家模型负责处理不同类型的输入数据。门控网络负责根据输入数据的特征，动态地分配每个专家模型的权重，从而得到最终的混合输出。

class MultiHeadAttention(hk.Module): def __init__( self, num_q_heads: int, num_kv_heads: int, key_size: int, *, with_bias: bool = True, value_size: Optional[int] = None, model_size: Optional[int] = None, attn_output_multiplier: 1.0, data_axis: Union[str, Tuple[str, ...]] = "data", model_axis: Union[str, Tuple[str, ...]] = "model", name: Optional[str] = None, ): super().__init__(name=name) self.num_q_heads = num_q_heads self.num_kv_heads = num_kv_heads self.key_size = key_size self.value_size = value_size or key_size self.model_size = model_size or key_size * num_q_heads self.data_axis = data_axis self.model_axis = model_axis self.attn_output_multiplier = attn_output_multiplier self.with_bias = with_bias

其grok是一个48头的多头注意力机制模型，其基于transformer的解码器部分，可见transformer模型框架的强大。目前所知的LLM的语言模型，比如GPT, Gemini，grok等模型都是基于transformer的模型。

https://x.ai/bloggithub.com/xai-org/grok.