fix tensor dimensions are not compatible for FP8 issue in sft (#8787)

Co-authored-by: yuhuang <[email protected]> Co-authored-by: Eric Harper <[email protected]>
NVIDIA · Jan 18, 2025 · 4df3fe5 · 4df3fe5
1 parent 8bf5144
commit 4df3fe5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/data/language_modeling/megatron/gpt_sft_chat_dataset.py b/nemo/collections/nlp/data/language_modeling/megatron/gpt_sft_chat_dataset.py
@@ -373,7 +373,7 @@ def collate_fn(self, batch):
         if self.pad_to_max_length:
             max_length = self.max_seq_length
         else:
-            max_length = min(self.max_seq_length, self._ceil_to_nearest(max_length, 8))
+            max_length = min(self.max_seq_length, self._ceil_to_nearest(max_length, 16))
         assert max_length <= self.max_seq_length
 
         if not self.get_attention_mask_from_fusion: