`gpt2_lm.generate()` always recompiles

[mattdangerw]

In our current implementation, every call to generate will trigger compilation, meaning generate() is quite slow.

There are actually a few reasons for this.

• Our sampler objects are regenerated each call generate().
• Inside the sampler object, sample() is explicitly recompiled each call.
• The sample object has python objects passed to it (a python function), meaning it will always retrace.

We need to make sure that subsequent calls to generate() with the same settings do not trigger extra compilation.

[mattdangerw]

The speedups of properly avoid recompilation are very large. This diff shows some hacky changes that avoid recompilation.

Given this test script:

import time
import keras_nlp

gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset("gpt2_base_en")
gpt2_lm.jit_compile = True
# First time to compile.
gpt2_lm.generate(["that's weird", "that's even weirder"], sampler="greedy", max_length=256)
# Measure subsequent runs.
start = time.time()
for i in range(25):
    gpt2_lm.generate(["that's weird", "that's even weirder"], sampler="greedy", max_length=256)
end = time.time()
print(end - start)

Time goes from 225s -> 52s on a V100 GPU.

[mattdangerw]

We also have an additional challenge with compilation, which is that our cache is always seeded outside the compiled function, in GTP2CausalLM. This means we are doing an entire forward pass eagerly, and don't see the benefits of compile.

As a whole, I am fairly sure we need to move compilation out of the sampler interface and into the task object. This will fit well with things like train_step and predict_step, where the model can own a compiled function, and only recompile when necessary (e.g. sampling strategy has changed).

[mattdangerw]

This is fixed!