text-generation-inference/server/text_generation_server/models
Nicolas Patry abd58ff82c
feat(server): Rework model loading (#344)
# What does this PR do?

Reworked the loading logic. Idea is to use cleaner loading code:

- Remove need for `no_init_weights`
- Remove all weird `bnb_linear` and `load_weights` and
`post_load_weights`.

New code layout:

- New class `Weights` in charge of handling loading the weights from
multiple files into appropiate tensors (potentially sharded)
- TP layers now are "shells", they contain the code to know what kind of
sharding we need + eventual `all_reduce`. They do not inherit from
linear, but they contain some kind of Linear instead
- the contained linear can be either FastLinear, BnbLinear or GPTq
Linear next.
- All modeling code is explictly made for sharding, process group is
just no-ops for non sharded code (removes a lot of test cases)

![Screenshot from 2023-05-19
23-19-59](https://github.com/huggingface/text-generation-inference/assets/204321/9a802654-74a3-488c-87a8-073743a6143f)

---------

Co-authored-by: Ubuntu <ubuntu@ip-172-31-41-161.taildb5d.ts.net>
Co-authored-by: Ubuntu <ubuntu@ip-172-31-41-161.ec2.internal>
Co-authored-by: OlivierDehaene <olivier@huggingface.co>
Co-authored-by: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
2023-06-08 14:51:52 +02:00
..
custom_modeling feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
__init__.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
bloom.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
causal_lm.py feat(server): only compute prefill logprobs when asked (#406) 2023-06-02 17:12:30 +02:00
flash_causal_lm.py feat(server): batch tokenization for flash causal lm (#411) 2023-06-05 16:09:41 +02:00
flash_llama.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
flash_neox.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
flash_rw.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
flash_santacoder.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
galactica.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
gpt_neox.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
model.py fix(server): fix has_position_ids (#395) 2023-06-01 11:41:35 +02:00
opt.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
rw.py feat(server): add retry on download (#384) 2023-05-31 10:57:53 +02:00
santacoder.py feat(server): support trust_remote_code (#363) 2023-05-23 20:40:39 +02:00
seq2seq_lm.py feat(server): only compute prefill logprobs when asked (#406) 2023-06-02 17:12:30 +02:00
t5.py feat(server): Rework model loading (#344) 2023-06-08 14:51:52 +02:00
types.py feat(server): support vectorized warpers in flash causal lm (#317) 2023-05-26 12:30:27 +02:00