AI-Hypercomputer / AI-Hypercomputer/JetStream

Refactor jestream to allow different tokenizers

Offen
#45 1 Kommentar 0 Reaktionen 1 zugewiesene Person Beansprucht von @bhavya01 Auf GitHub ansehen
Vorherrschende Sprache
Python
Sterne
457
Forks
67
PR-Merge-Kennzahlen
Keine gemergten PRs in 30 T.

Beschreibung

## 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.

Beitragsleitfaden

Beitragsleitfaden öffnen

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.