Implement parallel Inference and Generation

[Nintorac]

note: commits fix overlap should be checked in reverse and feat: add stopword checker + iterable generate function are from the implement stopwords PR.

I've only tested with the strategy cpu fp32 *1 but I think it should work for all strategies.

The parallel sampling method builds heavily on logic implemented in the stopword PR and the commits don't really define good boundaries. Let me know if you want me to fix that. If you merge this then the other can deleted

Otherwise seems to be running quite smoothly with what I've tried. Let me know what you think about it all and if there is any obvious math issues I've introduced!

I've written a bit on #dev-chatrwkv on discord about this

---

some examples that I should document

initialise

import gc
import logging
import os
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
os.environ["RWKV_JIT_ON"] = "0"
os.environ["RWKV_CUDA_ON"] = "0"
from rwkv.model import RWKV  # noqa: E402
from rwkv.utils import PIPELINE, PIPELINE_ARGS  # noqa: E402

ctx_limit = 4096

models = {
    "raven-14b-ctx4096": {
        "repo_id": "BlinkDL/rwkv-4-raven",
        "title": "RWKV-4-Raven-14B-v8-Eng-20230408-ctx4096",
    },
    "raven-7b-ctx4096": {
        "repo_id": "BlinkDL/rwkv-4-raven",
        "title": "RWKV-4-Raven-7B-v7-Eng-20230404-ctx4096",
    },
    "raven-7b-ctx1024": {
        "repo_id": "BlinkDL/rwkv-4-pile-7b",
        "title": "RWKV-4-Pile-7B-Instruct-test4-20230326",
    },
    "rwkv-4-pile-169m": {
        "repo_id": "BlinkDL/rwkv-4-pile-169m",
        "title": "RWKV-4-Pile-169M-20220807-8023",
    },
        "raven-1b-ctx4096": {
        "repo_id": "BlinkDL/rwkv-4-raven",
        "title": "RWKV-4-Raven-1B5-v11-Eng99%-Other1%-20230425-ctx4096",
    },
}

model = "raven-1b-ctx4096"
model_params = models[model]


def fetch_tokenizer(tokenizer_path: Path):
    url = "https://huggingface.co/spaces/BlinkDL/Raven-RWKV-7B/raw/main/20B_tokenizer.json"
    tokenizer_path.parent.mkdir(exist_ok=True)

    response = requests.get(url)
    tokenizer_path.write_bytes(response.content)


def get_model():
    tokenizer_path = Path(__file__).parent / "20B_tokenizer.json"
    if not tokenizer_path.exists():
        fetch_tokenizer(tokenizer_path)

    model_path = hf_hub_download(
        repo_id=model_params["repo_id"], filename=f"{model_params['title']}.pth"
    )

    model = RWKV(
        model=model_path, strategy="cpu fp32 *0"
    )  # stream mode w/some static

    pipeline = PIPELINE(model, str(tokenizer_path))

    return model, pipeline

model, pipeline = get_model()

Generate parallel

>>> questions = [
    'Bob: whats the day?\n\nAlice:',
    'Bob: whats the time?\n\nAlice:',
    'Bob: this is an obnoxiosuly long question, so much longer than the rest?\n\nAlice:',
    'Bob: whats the sound?\n\nAlice:',
    'Bob: whats the delay?\n\nAlice:',
]
>>> for i in pipeline.igenerate(questions, token_count=250, args=PIPELINE_ARGS(chunk_len=256, stop_words='Bob:')):
>>>     print(i)

Generate single

>>> for i in pipeline.igenerate('Bob: whats the day?\n\nAlice:', args=PIPELINE_ARGS(stop_words='Bob:')):
>>>     print(i, sep='', end='', flush=True)
 It is Saturday, April 28th.

Infer single

>>> state = pipeline.infer('Bob: whats the day?\n\nalice:')
[tensor([[ 0.0067, -0.0483, -0.1088,  ..., -0.0043,  0.0065,  0.0251]])]

Infer parallel

>>> pipeline.infer(questions)
[tensor([[ 0.0067, -0.0483, -0.1088,  ..., -0.0043,  0.0065,  0.0251],
         [-0.8015, -0.3936,  2.3798,  ...,  0.3735, -0.0855, -0.8532],
         [ 0.0090,  0.0298, -0.0651,  ...,  0.0129, -0.0053,  0.0163],
         [ 1.6057,  2.8717,  3.0602,  ...,  1.0000,  1.0000,  1.0000],
         [ 2.6282,  2.8281,  1.6604,  ..., -2.7151,  2.5806,  1.6659]]),
 tensor([[ 0.0067, -0.0483, -0.1088,  ..., -0.0043,  0.0065,  0.0251],
         [-0.7918, -0.8789,  1.8799,  ...,  0.3735, -0.0855, -0.8532],
         [ 0.0082,  0.0315, -0.0668,  ...,  0.0112, -0.0041,  0.0179],
         [ 1.5889,  2.3284,  2.8709,  ...,  1.0000,  1.0000,  1.0000],
         [ 2.6282,  2.8281,  1.6604,  ..., -2.7151,  2.5806,  1.6659]]),
 tensor([[ 6.6986e-03, -4.8346e-02, -1.0882e-01,  ..., -4.3214e-03,
           6.4733e-03,  2.5146e-02],
         [-7.9228e-01,  1.3874e+00,  2.3622e+00,  ...,  3.7351e-01,
          -8.5515e-02, -8.5321e-01],
         [ 8.2001e-03,  3.2023e-02, -6.6461e-02,  ...,  1.0814e-02,
          -2.8179e-03,  1.6461e-02],
         [ 1.5827e+00,  2.4019e+00,  2.9919e+00,  ...,  1.0000e+00,
           1.0000e+00,  1.0000e+00],
         [ 2.6282e+00,  3.2905e+00,  1.6604e+00,  ..., -2.7151e+00,
           2.5806e+00,  1.6659e+00]]),
 tensor([[ 6.6986e-03, -4.8346e-02, -1.0882e-01,  ..., -4.3214e-03,
           6.4733e-03,  2.5146e-02],
         [-7.8120e-01, -1.3236e+00,  1.9415e+00,  ...,  3.7351e-01,
          -8.5515e-02, -8.5321e-01],
         [ 1.1746e-02,  3.3171e-02, -6.4305e-02,  ...,  1.0538e-02,
          -5.1670e-04,  1.7817e-02],
         [ 1.5925e+00,  2.5117e+00,  3.0401e+00,  ...,  1.0000e+00,
           1.0000e+00,  1.0000e+00],
         [ 2.6282e+00,  2.8281e+00,  1.6604e+00,  ..., -2.7151e+00,
           2.5806e+00,  1.6659e+00]]),
 tensor([[ 6.6986e-03, -4.8346e-02, -1.0882e-01,  ..., -4.3214e-03,
           6.4733e-03,  2.5146e-02],
         [-1.1770e-02, -4.1843e-01,  2.1849e+00,  ...,  3.7351e-01,
          -8.5515e-02, -8.5321e-01],
         [ 1.4248e-02,  4.4202e-02, -7.8099e-02,  ...,  1.5142e-02,
          -5.2775e-04,  2.7968e-02],
         [ 3.9415e+00,  2.2419e+00,  7.9089e+00,  ...,  1.0000e+00,
           1.0000e+00,  1.0000e+00],
         [ 2.3895e+00,  3.5548e+00,  1.8283e+00,  ..., -2.7151e+00,
           2.5806e+00,  1.6659e+00]])]