add length pad for flex attn

Signed-off-by: h-guo18 <[email protected]>

Author: h-guo18

examples/speculative_decoding/eagle_utils.py

@@ -236,7 +236,10 @@ def __getitem__(self, i) -> dict[str, torch.Tensor]:
 
 
 def make_eagle_supervised_data_module(
-    tokenizer: transformers.PreTrainedTokenizer, data_args, use_offline_training: bool
+    tokenizer: transformers.PreTrainedTokenizer,
+    data_args,
+    use_offline_training: bool,
+    pad_length=None,
 ) -> dict:
     """Make dataset and collator for supervised fine-tuning.
 
@@ -303,7 +306,7 @@ def make_eagle_supervised_data_module(
         train_dataset = dataset_cls(data_json[: int(len(data_json) * 0.95)], tokenizer=tokenizer)
         eval_dataset = dataset_cls(data_json[int(len(data_json) * 0.95) :], tokenizer=tokenizer)
 
-        data_collator = DataCollatorWithPadding()
+        data_collator = DataCollatorWithPadding(pad_length=pad_length)
 
     return {
         "train_dataset": train_dataset,
@@ -313,6 +316,9 @@ def make_eagle_supervised_data_module(
 
 
 class DataCollatorWithPadding:
+    def __init__(self, pad_length=None):
+        self.pad_length = pad_length
+
     def paddingtensor2d(self, intensors, length):
         n, dim = intensors.shape
         padding_tensor = torch.zeros(length - n, dim, dtype=intensors.dtype)
@@ -325,7 +331,11 @@ def paddingtensor(self, intensors, length):
         return outtensors
 
     def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]:
-        max_length = max(item["input_ids"].shape[0] for item in features)
+        max_length = (
+            self.pad_length
+            if self.pad_length is not None
+            else max(item["input_ids"].shape[0] for item in features)
+        )
         batch_input_ids = torch.stack(
             [self.paddingtensor(item["input_ids"], max_length) for item in features]
         )
@@ -357,7 +367,11 @@ def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]:
             raise ValueError("No kwargs found in batch features. Offline data required.")
 
         features = [item["kwargs"]["base_model_outputs"] for item in features]
-        max_hs_length = max(item["base_model_hidden_states"].shape[0] for item in features)
+        max_hs_length = (
+            max(item["base_model_hidden_states"].shape[0] for item in features)
+            if self.pad_length is None
+            else self.pad_length
+        )
 
         batch_hidden_states = torch.stack(
             [

examples/speculative_decoding/main.py

@@ -229,7 +229,9 @@ def train():
     if training_args.mode == "medusa":
         data_module = make_medusa_supervised_data_module(tokenizer, data_args)
     elif training_args.mode in ["eagle1", "eagle3"]:
-        data_module = make_eagle_supervised_data_module(tokenizer, data_args, use_offline_training)
+        data_module = make_eagle_supervised_data_module(
+            tokenizer, data_args, use_offline_training, pad_length=training_args.training_seq_len
+        )
 
     class ARValidationCallback(TrainerCallback):
         def __init__(self, ar_validate_steps: int = 500):

modelopt/torch/speculative/plugins/transformers.py

@@ -316,7 +316,7 @@ def forward(
                 hidden_states,
                 attention_mask=attention_mask,
                 position_ids=position_ids,
-                past_key_values=past_key_values,
+                past_key_value=past_key_values,
                 output_attentions=output_attentions,
                 use_cache=use_cache,
                 position_embeddings=position_embeddings,