Llama

[Ivan-Zhou]

Implement Llama based on HF implementation and the paper.

A few noteable difference from Gpt2:

• Rotary Positional Embedding;
• No dropout being used (therefore, many of the call() doesn't need key as input).

Tasks

• Add Llama code based on HF implementation
• Refactor with Haliax
• Match Levanter's state dict with HF's
• Roundtrip test
• Ensure perf roughly matches GPT-2 implementation at scale

[dlwh]

btw i fixed the problem in latest haliax. You should still make them static fields, but you won't get an error there

…

On Wed, Aug 16, 2023 at 7:58 AM Ivan Zhou ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In src/levanter/models/llama.py <#270 (comment)> : > + residual = x + x = self.ln_2(x) + mlp_output = self.mlp(x) + output = residual + mlp_output + return output + + +class LlamaTransformer(StateDictSerializationMixin, eqx.Module): + config: LlamaConfig = eqx.static_field() + layers: Stacked[LlamaDecoderLayer] + ln_f: hnn.LayerNorm + + @staticmethod + def init(config: LlamaConfig, *, key) -> "LlamaTransformer": + # TODO: here it reports an error that is related to _get_rotary_emb() in LlamaAttention + # TypeError: Output from batched function Axis(name='head_size', size=4) with type <class 'haliax.axis.Axis'> is not a valid JAX type Will make a try — Reply to this email directly, view it on GitHub <#270 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACLIKWRMRODVFPZROCRK3XVTN3BANCNFSM6AAAAAA3BRV53M> . You are receiving this because you commented.Message ID: ***@***.***>

[Ivan-Zhou]

Untie word_embeddings at LMHead is done. The only pending issue is Jax leakage when loading HF weight to Levanter's model.

[Ivan-Zhou]

Great thanks to @dlwh for helping with the roundtrip tests and massively improve the code style & taste in this PR 👍