Speculative sampling

[haim-barad]

Works with DollyV2 for machines with at least 64GB local memory. By default, we use GPT2 for smaller machines.

Leveraging speculative sampling with KV caching, this notebook's code will generate text using standard autoregressive sampling and speculative sampling and compare the times needed to generate N tokens. By default, N is set to 100. Another parameter K is set to 5 but can be adjusted to determine the number of candidate tokens to generate from the smaller draft model.

This code uses gradio to provide an interface labeling the models used and handling the input.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[andrei-kochin]

@haim-barad please check the CI failures as it seems some of the are quite easy to fix

[apaniukov]

Line #31.            res = target_model(x_draft_target_input, attention_mask=torch.ones(x_draft_target_input.size(), dtype=torch.long), use_cache=False)

Why use_cache=False? You should use target_past_kv for target_model inference based on x_draft_target_input:

if target_past_kv is None:
x_draft_target_input = torch.cat((x, x_draft), dim=1)

else:

x_draft_target_input = x_draft

Reply via ReviewNB

[apaniukov]

Line #37.                q_cum = torch.cat((q_cum, q), dim=1)

q_cum is not used anywhere.

Reply via ReviewNB

[apaniukov]

Line #46.                if np.random.random() < min(1, (q_item / p_item)):  # accepted

q_item is logits, p_item is probability after softmax. You should apply softmax to q first.

Reply via ReviewNB

[apaniukov]

Line #58.            target_past_kv = target_new_past_kv

You should update target_past_kv based on the number of accepted tokens.

Reply via ReviewNB

[haim-barad]

Code_check failure is another notebook. Spelling errors are ignoring the tag (as noted in your instructions). Treon (ubuntu for 3.8 and 3.9) fails due to timeout for GPU... but I clearly have CPU as a default device.

[eaidova]

@haim-barad, for spell check you need to add unknown words reported by tools into vocabulary. You can find more info about that contributing guide

[haim-barad]

I had added the required entries and now I see the latest push created some conflicts, I apologize for this, but can someone with permissions do a merge?

[haim-barad]

code check is failing for 263-latent-consistency-models-image-generation, not this notebook.

[haim-barad]

Both code_check and docker_treon are failing on other notebooks, not because of this PR. Please approve. Thanks.

[haim-barad]

There's still a small link to fix. I see this was renumbered to 266. There's a "open in Collab" button in the readme - the link needs to be 266 instead of 265. Can one of you fix it to avoid the review process?