testing correctness between kv cache, non kvcache, and eager

Differential Revision: D55465401

Author: JacobSzwejbka

examples/models/llama2/TARGETS

@@ -57,7 +57,7 @@ runtime.python_binary(
     deps = [
         ":export_library",
         "//caffe2:torch",
-        "//executorch/extension/pybindings:aten_lib",
+        "//executorch/extension/pybindings:portable_lib",
     ],
 )
 

examples/models/llama2/export_kv_d32.txt

@@ -26,6 +26,9 @@
 
 from executorch.examples.models.llama2.llama_transformer import Transformer
 from executorch.exir.backend.backend_details import CompileSpec
+from executorch.extension.pybindings.portable_lib import (
+    _load_for_executorch_from_buffer,
+)
 
 from executorch.sdk.etrecord import generate_etrecord
 from executorch.util.activation_memory_profiler import generate_memory_trace
@@ -774,6 +777,33 @@ def _export_llama(modelname, args) -> str:  # noqa: C901
     else:
         output_file = f"{builder.output_dir}/{modelname}.pte"
 
+    builder.export_program.dump_executorch_program(True)
+    et_model = _load_for_executorch_from_buffer(builder.export_program.buffer)
+    eager_model = builder.model
+    input1 = torch.tensor([[0]], dtype=torch.long)
+    input2 = torch.tensor([0], dtype=torch.long)
+    if args.use_kv_cache:
+        for i in range(0, 100):
+            eager_res = eager_model(input1, input2)
+            et_res = et_model((input1, input2))
+            assert len(et_res) == 1
+            print("eager res kv", i, eager_res)
+            print("et res kv", i, et_res[0])
+            # assert torch.allclose(eager_res, et_res[0], atol=1e-05, rtol=1e-05)
+            input1 = torch.tensor([[i + 1]], dtype=torch.long)
+            input2 = torch.tensor([i + 1], dtype=torch.long)
+    else:
+        for i in range(0, 100):
+            eager_res = eager_model(input1)
+            et_res = et_model((input1,))
+            assert len(et_res) == 1
+            print("eager res", i, eager_res)
+            print("et res", i, et_res[0])
+            # assert torch.allclose(eager_res, et_res[0], atol=1e-05, rtol=1e-05)
+            input1 = torch.cat(
+                [input1, torch.tensor([[i + 1]], dtype=torch.long)], dim=1
+            )
+
     builder.save_to_pte(output_file)
 
     return output_file