Code

Add inference cache (cleaner version)#4

Code

Open

borzunov wants to merge 19 commits intoweight-sharinglearning-at-home/dalle-pytorch:weight-sharingfrom

inference-cache-v2learning-at-home/dalle-pytorch:inference-cache-v2Copy head branch name to clipboard

Commits

Commits on Dec 20, 2021

• Add FFN caching

  borzunov

  committed
  1cd8e20

  Copy full SHA for 1cd8e20

• Cache to_qkv and to_out in sparse attn, add debug prints

  borzunov

  committed
  4d431ac

  Copy full SHA for 4d431ac

• Cache full Attention

  borzunov

  committed
  2603776

  Copy full SHA for 2603776

• Remove debug outputs

  borzunov

  committed
  112ea05

  Copy full SHA for 112ea05

• Cache pre-logits MLP

  borzunov

  committed
  6ba4cb6

  Copy full SHA for 6ba4cb6

• Further optimize attention caching

  borzunov

  committed
  c333ea7

  Copy full SHA for c333ea7

Commits on Dec 21, 2021

• Fix mask in attention

  borzunov

  committed
  1fd45ca

  Copy full SHA for 1fd45ca

• Don't cache MLPs since we can just pass only last item

  borzunov

  committed
  2b77018

  Copy full SHA for 2b77018

• Revert excess changes in attentions

  borzunov

  committed
  8e8dea8

  Copy full SHA for 8e8dea8

• Rename FixCacheKey -> CachedAs

  borzunov

  committed
  df89951

  Copy full SHA for df89951

• Save the current offset in cache

  borzunov

  committed
  059fe1b

  Copy full SHA for 059fe1b

• Use static masks to simulate axial attn

  borzunov

  committed
  b76b78e

  Copy full SHA for b76b78e

• Add option to disable caching

  borzunov

  committed
  4c833a2

  Copy full SHA for 4c833a2

• Make the cached version work

  borzunov

  committed
  1ff47c6

  Copy full SHA for 1ff47c6

• Speed up PreShiftToken

  borzunov

  committed
  94fda36

  Copy full SHA for 94fda36

Commits on Jan 10, 2022

• Remove excess changes

  borzunov

  committed
  adfce34

  Copy full SHA for adfce34

• Add NonCached wrapper

  borzunov

  committed
  59cfc49

  Copy full SHA for 59cfc49

• Rename use_static_masks -> optimize_for_inference

  borzunov

  committed
  4f496a4

  Copy full SHA for 4f496a4

• Improve names and comments

  borzunov

  committed
  732226d

  Copy full SHA for 732226d