Code

Add Llama Flax Implementation#24587

Code

Merged

ArthurZucker merged 91 commits intohuggingface:mainhuggingface/transformers:mainfrom vvvm23:add-llama-flaxvvvm23/transformers:add-llama-flaxCopy head branch name to clipboardDec 7, 2023

+1,071-16

Commits

Commits on Nov 5, 2023

• Copies modeling_flax_gpt_neo.py to start

  vvvm23

  committed
  415f0d2

  Copy full SHA for 415f0d2

• MLP Block. WIP Attention and Block

  vvvm23

  committed
  76a599c

  Copy full SHA for 76a599c

• Adds Flax implementation of LlamaMLP

  Show description for 4f5654d
  vvvm23

  committed
  4f5654d

  Copy full SHA for 4f5654d

• Adds FlaxLlamaRMSNorm layer

  Show description for 36c48fa
  vvvm23

  committed
  36c48fa

  Copy full SHA for 36c48fa

• Adds FlaxLlamaAttention

  Show description for a38b097
  vvvm23

  committed
  a38b097

  Copy full SHA for a38b097

• Adds FlaxLlamaDecoderLayer

  Show description for 5738666
  vvvm23

  committed
  5738666

  Copy full SHA for 5738666

• debugging rotary mismatch

  Show description for 401a72f
  vvvm23

  committed
  401a72f

  Copy full SHA for 401a72f

• fixes bug with decoder layer

  Show description for 5177bfd
  vvvm23

  committed
  5177bfd

  Copy full SHA for 5177bfd

• adds markers for what to implement next

  Show description for 578e0d9
  vvvm23

  committed
  578e0d9

  Copy full SHA for 578e0d9

• implements FlaxLlamaBlockCollection]

  Show description for 9f74d83
  vvvm23

  committed
  9f74d83

  Copy full SHA for 9f74d83

• Adds FlaxLlamaModule

  Show description for a60b00f
  vvvm23

  committed
  a60b00f

  Copy full SHA for a60b00f

• adds FlaxLlamaForCausalLMModule

  Show description for c7ac55b
  vvvm23

  committed
  c7ac55b

  Copy full SHA for c7ac55b

• start porting pretrained wrappers

  Show description for b6dff5a
  vvvm23

  committed
  b6dff5a

  Copy full SHA for b6dff5a

• cleanup, quality, style

  vvvm23

  committed
  7bf3567

  Copy full SHA for 7bf3567

• readds return_dict and model output named tuples

  vvvm23

  committed
  99d40a0

  Copy full SHA for 99d40a0

• (tentatively) pretrained wrappers work 🔥

  vvvm23

  committed
  4eebae9

  Copy full SHA for 4eebae9

• fixes numerical mismatch in FlaxLlamaRMSNorm

  Show description for 4bb4206
  vvvm23

  committed
  4bb4206

  Copy full SHA for 4bb4206

• [WIP] debugging numerics

  vvvm23

  committed
  b78671e

  Copy full SHA for b78671e

• numerical match

  Show description for 3f7bc54
  vvvm23

  committed
  3f7bc54

  Copy full SHA for 3f7bc54

• adds in model and integration tests for Flax Llama

  Show description for e386c85
  vvvm23

  committed
  e386c85

  Copy full SHA for e386c85

• adds missing TYPE_CHECKING import and make fixup

  vvvm23

  committed
  539d041

  Copy full SHA for 539d041

• adds back missing docstrings

  Show description for c695d0a
  vvvm23

  committed
  c695d0a

  Copy full SHA for c695d0a

• commenting out equivalence test as can just use common

  vvvm23

  committed
  1776e75

  Copy full SHA for 1776e75

• debugging

  vvvm23

  committed
  3b4f55a

  Copy full SHA for 3b4f55a

• Fixes bug where mask and pos_ids were swapped in pretrained models

  Show description for 3fa9f2e
  vvvm23

  committed
  3fa9f2e

  Copy full SHA for 3fa9f2e

• cleanup of modeling file

  vvvm23

  committed
  9d4bdad

  Copy full SHA for 9d4bdad

• cleanup of test file

  vvvm23

  committed
  ffb5e47

  Copy full SHA for ffb5e47

• Resolving simpler review comments

  vvvm23

  committed
  020bd4e

  Copy full SHA for 020bd4e

• addresses more minor review comments

  vvvm23

  committed
  e9e391f

  Copy full SHA for e9e391f

• fixing introduced pytest errors from review

  vvvm23

  committed
  ff0818f

  Copy full SHA for ff0818f

• wip additional slow tests

  vvvm23

  committed
  d18daad

  Copy full SHA for d18daad

• wip tests

  Show description for 230abeb
  vvvm23

  committed
  230abeb

  Copy full SHA for 230abeb

• make quality, make style

  vvvm23

  committed
  b19213d

  Copy full SHA for b19213d

• adds slow integration tests

  Show description for 2959abd
  vvvm23

  committed
  2959abd

  Copy full SHA for 2959abd

• make fix-copies

  vvvm23

  committed
  a5b587b

  Copy full SHA for a5b587b

• fix mangled function following make fix-copies

  vvvm23

  committed
  852e5e3

  Copy full SHA for 852e5e3

• adds missing type checking imports

  vvvm23

  committed
  fd85d5a

  Copy full SHA for fd85d5a

• fixes missing parameter checkpoint warning

  vvvm23

  committed
  fe5aed2

  Copy full SHA for fe5aed2

• more finegrained 'Copied from' tags

  Show description for 57b47c6
  vvvm23

  committed
  57b47c6

  Copy full SHA for 57b47c6

• swaps import guards

  Show description for b768559
  vvvm23

  committed
  b768559

  Copy full SHA for b768559

• removing inv_freq again as pytorch version has now removed

  vvvm23

  committed
  ac3f74f

  Copy full SHA for ac3f74f

• attempting to get CI to pass

  vvvm23

  committed
  05cade4

  Copy full SHA for 05cade4

• adds doc entries for llama flax models

  vvvm23

  committed
  3bf0b8b

  Copy full SHA for 3bf0b8b

• fixes typo in __init__.py imports

  vvvm23

  committed
  211a72b

  Copy full SHA for 211a72b

• adds back special equivalence tests

  Show description for 27a7522
  vvvm23

  committed
  27a7522

  Copy full SHA for 27a7522

• overrides tests with dummy to see if CI passes

  Show description for 67f300c
  vvvm23

  committed
  67f300c

  Copy full SHA for 67f300c

• adds my contribution to docs

  vvvm23

  committed
  2ec5c20

  Copy full SHA for 2ec5c20

• make style; make quality

  vvvm23

  committed
  609a113

  Copy full SHA for 609a113

• replaces random masking with fixed to work with flax version

  vvvm23

  committed
  224f546

  Copy full SHA for 224f546

• make quality; make style

  vvvm23

  committed
  7de8b58

  Copy full SHA for 7de8b58

• Update src/transformers/models/llama/modeling_flax_llama.py

  Show description for 20b5767
  

  vvvm23

  and

  sanchit-gandhi

  committed
  20b5767

  Copy full SHA for 20b5767

• Update src/transformers/models/llama/modeling_flax_llama.py

  Show description for ac4183c
  

  vvvm23

  and

  sanchit-gandhi

  committed
  ac4183c

  Copy full SHA for ac4183c

• Update src/transformers/models/llama/modeling_flax_llama.py

  Show description for bd5451a
  

  vvvm23

  and

  sanchit-gandhi

  committed
  bd5451a

  Copy full SHA for bd5451a

• Update src/transformers/models/llama/modeling_flax_llama.py

  Show description for c997a38
  

  vvvm23

  and

  sanchit-gandhi

  committed
  c997a38

  Copy full SHA for c997a38

• Update src/transformers/models/llama/modeling_flax_llama.py

  Show description for 5019d4c
  

  vvvm23

  and

  sanchit-gandhi

  committed
  5019d4c

  Copy full SHA for 5019d4c

• Update src/transformers/models/llama/modeling_flax_llama.py

  Show description for 4df7730
  

  vvvm23

  and

  sanchit-gandhi

  committed
  4df7730

  Copy full SHA for 4df7730

• updates x->tensor in rotate_half

  vvvm23

  committed
  f8ccb05

  Copy full SHA for f8ccb05

• addresses smaller review comments

  vvvm23

  committed
  b01cb70

  Copy full SHA for b01cb70

• Update docs/source/en/model_doc/llama.md

  Show description for 6848c63
  

  vvvm23

  and

  sanchit-gandhi

  committed
  6848c63

  Copy full SHA for 6848c63

• adds integration test class

  vvvm23

  committed
  9994b91

  Copy full SHA for 9994b91

• adds dtype to rotary embedding to cast outputs

  vvvm23

  committed
  d248925

  Copy full SHA for d248925

• adds type to flax llama rotary layer

  vvvm23

  committed
  f1fc40a

  Copy full SHA for f1fc40a

• make style

  vvvm23

  committed
  1f7cb9b

  Copy full SHA for 1f7cb9b

• make fix-copies

  vvvm23

  committed
  be7be91

  Copy full SHA for be7be91

• Apply suggestions from code review

  Show description for 3a7a3ae
  

  vvvm23

  and

  sanchit-gandhi

  committed
  3a7a3ae

  Copy full SHA for 3a7a3ae

• applies suggestions from review

  vvvm23

  committed
  20e6b35

  Copy full SHA for 20e6b35

• Update modeling_flax_llama.py

  vvvm23

  committed
  3da6a6a

  Copy full SHA for 3da6a6a

• make fix-copies

  vvvm23

  committed
  5f5ca1d

  Copy full SHA for 5f5ca1d

• Update tests/models/llama/test_modeling_llama.py

  Show description for 9166130
  

  vvvm23

  and

  sanchit-gandhi

  committed
  9166130

  Copy full SHA for 9166130

• Update src/transformers/models/llama/modeling_flax_llama.py

  Show description for d8570b4
  

  vvvm23

  and

  sanchit-gandhi

  committed
  d8570b4

  Copy full SHA for d8570b4

• fixes shape mismatch in FlaxLlamaMLP

  vvvm23

  committed
  6d7a930

  Copy full SHA for 6d7a930

• applies some suggestions from reviews

  vvvm23

  committed
  39b55f8

  Copy full SHA for 39b55f8

• casts attn output logits to f32 regardless of dtype

  vvvm23

  committed
  a6d8c06

  Copy full SHA for a6d8c06

• adds attn bias using LlamaConfig.attention_bias

  vvvm23

  committed
  9718fed

  Copy full SHA for 9718fed

• adds Copied From comments to Flax Llama test

  vvvm23

  committed
  fc6554b

  Copy full SHA for fc6554b

• mistral and persimmon test change -copy from llama

  vvvm23

  committed
  f9fd7b6

  Copy full SHA for f9fd7b6

• updates docs index

  vvvm23

  committed
  8b1f374

  Copy full SHA for 8b1f374

• removes Copied from in tests

  Show description for d9c7af6
  vvvm23

  committed
  d9c7af6

  Copy full SHA for d9c7af6

• quality and style

  vvvm23

  committed
  30fe8e1

  Copy full SHA for 30fe8e1

Commits on Nov 9, 2023

• ignores FlaxLlama input docstring

  vvvm23

  committed
  d21d306

  Copy full SHA for d21d306

• Merge branch 'main' into add-llama-flax

  vvvm23

  committed
  d14d339

  Copy full SHA for d14d339

Commits on Nov 24, 2023

• Merge branch 'main' into add-llama-flax

  vvvm23

  committed
  2b9e410

  Copy full SHA for 2b9e410

• adds revision to _CHECKPOINT_FOR_DOC

  vvvm23

  committed
  3390651

  Copy full SHA for 3390651

• repo consistency and quality

  vvvm23

  committed
  f303188

  Copy full SHA for f303188

• removes unused import

  vvvm23

  committed
  aea3e03

  Copy full SHA for aea3e03

• removes copied from from Phi test

  Show description for 842b550
  vvvm23

  committed
  842b550

  Copy full SHA for 842b550

Commits on Dec 4, 2023

• Merge branch 'main' into add-llama-flax

  vvvm23

  committed
  7e69179

  Copy full SHA for 7e69179

• adds _REAL_CHECKPOINT_FOR_DOC

  vvvm23

  committed
  193052f

  Copy full SHA for 193052f

Commits on Dec 5, 2023

• removes refs from pr tests

  vvvm23

  committed
  b33a54d

  Copy full SHA for b33a54d

• reformat to make ruff happy

  vvvm23

  committed
  b9ed34f

  Copy full SHA for b9ed34f

• Merge branch 'main' into add-llama-flax

  vvvm23

  committed
  a15b844

  Copy full SHA for a15b844