set seed for mlm convergence

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

recipes/esm2_native_te_nvfsdp_thd/test_thd_format.py

@@ -52,6 +52,7 @@ def test_thd_format():
         DataCollatorForLanguageModeling(
             tokenizer=tokenizer,
             mlm_probability=0.15,
+            seed=42,
         ),
         DataCollatorWithFlattening(
             return_flash_attn_kwargs=True,
@@ -137,7 +138,7 @@ def test_thd_format_with_different_batch_sizes():
     single_batch = [{"input_ids": [0, 5, 10, 15, 1], "attention_mask": [1, 1, 1, 1, 1]}]
 
     data_collator = MLMDataCollatorWithFlattening(
-        DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15),
+        DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15, seed=42),
         DataCollatorWithFlattening(return_flash_attn_kwargs=True),
     )
 
@@ -156,7 +157,9 @@ def test_thd_format_sequence_lengths():
     original_lengths = [len(seq["input_ids"]) for seq in batch]
 
     data_collator = MLMDataCollatorWithFlattening(
-        DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.0),  # No masking for length test
+        DataCollatorForLanguageModeling(
+            tokenizer=tokenizer, mlm_probability=0.0, seed=42
+        ),  # No masking for length test
         DataCollatorWithFlattening(return_flash_attn_kwargs=True),
     )
 
@@ -174,7 +177,7 @@ def test_thd_format_tensor_types():
     batch = create_test_batch()
 
     data_collator = MLMDataCollatorWithFlattening(
-        DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15),
+        DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15, seed=42),
         DataCollatorWithFlattening(return_flash_attn_kwargs=True),
     )
 
@@ -208,7 +211,7 @@ def test_mlm_data_collator_integration():
     # Test with different MLM probabilities
     for mlm_prob in [0.0, 0.15, 0.3]:
         data_collator = MLMDataCollatorWithFlattening(
-            DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=mlm_prob),
+            DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=mlm_prob, seed=42),
             DataCollatorWithFlattening(return_flash_attn_kwargs=True),
         )