1/n flamingo e2e ppl

Author: Gasoonjia

.gitignore

@@ -30,3 +30,6 @@ system_info.txt
 # build artifacts
 checkpoints/
 exportedModels/
+
+# test script
+_torchchat_test_script.py

torchchat/generate.py

@@ -21,7 +21,7 @@
 import torch._inductor.config
 
 try:
-    from _torchchat_test_script import flamingo_transform, padded_collate
+    from _torchchat_test_script import flamingo_transform
 except ImportError:
     pass
 
@@ -38,8 +38,9 @@
 from torchchat.utils.device_info import get_device_info
 
 # torchtune model definition dependencies
-from torchtune.data import Message
-from torchtune.generation._generation import sample as tune_sample
+from torchtune.data import Message, padded_collate_tiled_images_and_mask
+
+from torchtune.generation import sample as tune_sample
 from torchtune.models.llama3 import llama3_tokenizer
 from torchtune.training import set_default_dtype
 
@@ -357,15 +358,25 @@ def prefill(
 
         if batch is not None:
             # TODO: Verify sequential prefill works with multimodal models
-            logits = model(**batch)[:, -1]
-            return tune_sample(logits, 0, 500)
+            tokens = batch["tokens"]
+            if 'encoder_input' in tokens:
+                encoder_input = tokens['encoder_input']
+            else:
+                encoder_input = None
+            
+            mask = batch["causal_mask"][None, :seq_len]
+            input_pos = batch["input_pos"][None, :seq_len]
+            encoder_mask = batch["encoder_mask"]
+
+            logits = model(tokens=tokens, mask=mask, encoder_input=encoder_input, input_pos=input_post, encoder_mask=encoder_mask)[:, -1]
+            return tune_sample(logits, temperature=0, top_k=500)
         elif sequential_prefill:
             for i in range(width):
                 x_sliced, ip_sliced = x[:, i].view(-1, 1), input_pos[i].view(-1)
                 # logging.debug(f"<sliced> x: {x_sliced}, input_pos: {ip_sliced}")
                 logits = model(x_sliced, ip_sliced)  # (x[:, i], input_pos[i])
         elif self.model.config.model_type == ModelType.Flamingo:
-            logits = model(x)
+            assert False, "Flamingo requires batch"
         else:
             # input_pos: [B, S]
             logits = model(x, input_pos)
@@ -387,10 +398,10 @@ def decode_one_token(
         assert input_pos.shape[-1] == 1
         x = x.view(1, -1)
         if model.config.model_type == ModelType.Flamingo:
-            if batch is not None:
-                logits = model(x, encoder_mask=batch["encoder_mask"][:, -1:])
-            else:
-                logits = model(x)
+            assert batch is not None, "Flamingo requires batch"
+            mask = batch["causal_mask"][None, input_pos.item(), None, :]
+            encoder_mask = batch["encoder_mask"][:, -1:]
+            logits = model(x, encoder_mask=encoder_mask, mask=mask, input_pos=input_pos)[:, -1:]
         else:
             logits = model(x, input_pos)
         # print(f"x: {x},\n  input_pos: {input_pos}\n")
@@ -593,7 +604,7 @@ def generate(
                     self.is_torchtune_model
                     or self.model.config.model_type == ModelType.Flamingo
                 ):
-                    model.setup_caches(max_batch_size=1, dtype=self.dtype)
+                    model.setup_caches(batch_size=1, dtype=self.dtype, encoder_max_seq_len=6404, decoder_max_seq_len=T_new)
                 else:
                     model.setup_caches(max_batch_size=1, max_seq_length=max_seq_length)
                 if is_speculative and draft_model is not model:
@@ -743,8 +754,16 @@ def chat(
 
             transform = flamingo_transform(str(self.tokenizer_args.tokenizer_path))
             data = transform({"messages": messages}, inference=True)
-            batch = padded_collate([data], self.builder_args.device)
-            batch.pop("mask")
+            batch = padded_collate_tiled_images_and_mask([data], pad_direction="left", pad_max_images=1)
+            seq_len = len(data["tokens"])
+            total_response_length = seq_len + generator_args.max_new_tokens
+            batch["causal_mask"] = torch.tril(
+                                        torch.ones(
+                                            size=(total_response_length, total_response_length),
+                                            dtype=torch.bool,
+                                        )
+                                    )
+            batch["encoder_mask"] = batch["encoder_mask"][:, :seq_len]
             encoded = batch["tokens"]
 
         else: