fix use seq_length instead of packing_buffer_size to set max number of tokens.

Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>

Author: sajadn

dfm/src/megatron/data/dit/diffusion_task_encoder_with_sp.py

@@ -65,7 +65,7 @@ def select_samples_to_pack(self, samples: List[DiffusionSample]) -> List[List[Di
         """
         Selects sequences to pack for mixed image-video training.
         """
-        results = first_fit_decreasing(samples, self.packing_buffer_size)
+        results = first_fit_decreasing(samples, self.seq_length)
         random.shuffle(results)
         return results
 

dfm/src/megatron/recipes/dit/dit.py

@@ -185,8 +185,8 @@ def pretrain_config(
         dataset=DiffusionDataModuleConfig(
             path=dataset_path,
             seq_length=2048,
-            task_encoder_seq_length=2048,
-            packing_buffer_size=8000,
+            task_encoder_seq_length=8000,
+            packing_buffer_size=32,
             micro_batch_size=micro_batch_size,
             global_batch_size=global_batch_size,
             num_workers=10,