AI-Hypercomputer / AI-Hypercomputer/JetStream

Refactor jestream to allow different tokenizers

Aberta
#45 1 comentário 0 reações 1 responsável Reivindicada por @bhavya01 Ver no GitHub
Linguagem predominante
Python
Estrelas
457
Forks
67
Métricas de merge de PRs
Nenhum PR com merge em 30d

Descrição

## Issue

Currently we assume few things in jetstream which hinders it's generalization:

1. tokenizer is SentencePiece based.
2. pad_id is 0
3. after encode, we pad to nearest power of 2
4. ResultToken itself is jax specific (the @struct.dataclass annotation requires it's Jax pytreeable).

These assumptions hinders generalization (i.e. support wider varieties of models).
Examples:
1. llama3 uses tiktoken instead of SentencePiece
2. llama3 uses pad_id of -1
3. Pytorch GPU does NOT need to pad to nearst power of 2.
4. Pytorch GPU version of jetstream would like to use `torch.Tensor` to hold the data, which is not jax-pytreeable.

## Proposal:

1. `EngineAPI.get_tokenizer` which returns the tokenizer, should be any object that implements the following interface:

```
def encode

def decode

@property
def pad_id

@property
def eos_id
```
Uses of tokenizer should restrict to only this methods.

In particular: `encode` should do both encoding and padding. So jetstream doesnt do any padding itself; the engine can choose how to pad (or not to pad) by returning a custom tokenizer object whose encode also does the padding.

2. Allow use different implementation for `ResultTokens`; same as `Prefix` and `DecodeState`. Implementations of the Engine can choose implementation of ResultTokens. jestream should interact with it only through it's 3 public methods (https://github.com/google/JetStream/blob/main/jetstream/engine/engine_api.py#L83)

```python
def get_result_at_slot
def convert_to_numpy
def copy_to_host_async
```
and are not allowed to access it's fields directly.

Guia de contribuição

Abrir o guia de contribuição

Avaliação

Esta issue ainda não foi avaliada.

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.