bigscience-workshop / bigscience-workshop/petals

More powerful session API

Open
#495 1 comment 2 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
10.6k
Forks
650
PR merge metrics
No merged PRs in 30d

Description

Based on a discussion on Discord between me and @borzunov in the webui discussion, I didn't want it to get lost.

So consider a simple program using sessions:
```Python
with model.inference_session(max_length=500) as sess:
output1 = model.generate('[Input 1]', max_new_tokens=2, session=sess, do_sample=True, temperature=0.9, top_p=0.6)
output2 = model.generate(`[Input 2]`, max_new_tokens=2, session=sess, do_sample=True, temperature=0.9, top_p=0.6)
```

Currently the way the session API works is that it keeps history, so the first command generates `[Output 1]` and then the second command generates starting from `[Input 1][Output 1][Input 2]`. But compared to the usual transformers API this is quite restrictive, it is really only useful for chat-like applications where you can never go back and edit anything.

It would be a more powerful API if instead the .generate() calls acted as though they were unrelated/independent, and then the session managed the reuse logic internally. So for example if you wanted the old behavior you would call `model.generate("[Input 1][Output 1][Input 2]")` in the second call, but if you didn't you could still do `model.generate("[Input 2]")`. It's fairly cheap to process a buffer of tokens in Python and analyze it for potential reuse patterns.

As far as the reuse logic, I have developed the outline for a little algorithm that I think will work in most cases.
```Python
def do_generation(old,new):
start =0
while(old[start] != new[0]):
start += 1
len = 0
while(old[start + len] == new[len]):
len += 1
reuse_inference(old[start:start+len])
reprocess_fresh(new[len:])
```

This supports three main use cases:
* prefix - similar to the old one, doing `a` then `ab` will reuse the blocks
* prefix with different suffix - doing `ab` then `ac` will reuse `a`
* rolling - doing `abc` then `bcd` will drop `a` then reuse `bc`, which can happen as you get long prompts that exceed the context length

Per Alexander B., this would actually be fairly easy to implement in petals, but currently, it is not yet implemented.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.