马斯克说若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等模型。
GSM8k:中学数学应用题MMLU:多学科多项选择题HumanEval:Python 代码完成任务数学:用 LaTeX 编写的初中和高中数学问题从开源代码可以看出,其模型是基于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.