add t5 --ci-config option

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
NVIDIA · Oct 18, 2024 · 862038c · 862038c
1 parent 448ff8c
commit 862038c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml
@@ -5141,13 +5141,15 @@ jobs:
         NVTE_FUSED_ATTN=0 NVTE_FLASH_ATTN=0 python tests/collections/llm/megatron_t5_pretraining.py \
         --devices=2 \
         --max-steps=3 \
+        --ci-config \
         --experiment-dir=tests/collections/llm/t5_pretrain_results/${{ github.run_id }} \
         --data-path=/home/TestData/nlp/megatron_t5/data/pile_val_small_bert_tokenizer_text_document \
         --index-mapping-dir=tests/collections/llm/t5_index_mappings/${{ github.run_id }}
 
         NVTE_FUSED_ATTN=0 NVTE_FLASH_ATTN=0 python tests/collections/llm/megatron_t5_pretraining.py \
         --devices=2 \
         --max-steps=6 \
+        --ci-config \
         --experiment-dir=tests/collections/llm/t5_pretrain_results/${{ github.run_id }} \
         --data-path=/home/TestData/nlp/megatron_t5/data/pile_val_small_bert_tokenizer_text_document \
         --index-mapping-dir=tests/collections/llm/t5_index_mappings/${{ github.run_id }}
@@ -5790,4 +5792,4 @@ jobs:
 
       - name: "Pipeline not successful, set exit code to 1"
         if: ${{ always() && steps.pipeline-conclusion.outputs.SUCCESS == 'false' }}
-        run: exit 1
+        run: exit 1
diff --git a/tests/collections/llm/megatron_t5_pretraining.py b/tests/collections/llm/megatron_t5_pretraining.py
@@ -42,6 +42,7 @@ def get_args():
     parser.add_argument('--data-path', type=str, help="Path to data file")
     parser.add_argument('--vocab-path', type=str, default=None, help="Path to vocab file")
     parser.add_argument('--index-mapping-dir', type=str, help="directory to write index mappings to")
+    parser.add_argument('--ci-config', action='store_true', help="small config for CI purposes")
 
     return parser.parse_args()
 
@@ -59,16 +60,16 @@ def get_args():
         paths=args.data_path,
         seq_length=512,
         seq_length_dec=128,
-        micro_batch_size=args.devices,
-        global_batch_size=2 * args.devices,
+        micro_batch_size=args.devices if args.ci_config else 64,
+        global_batch_size=2 * args.devices if args.ci_config else 512,
         seed=1234,
         tokenizer=tokenizer,
         split="99982,9,9",
         index_mapping_dir=args.index_mapping_dir,
     )
     t5_config = llm.t5.model.t5.T5Config(
-        num_layers=args.devices,
-        encoder_num_layers=args.devices,
+        num_layers=args.devices if args.ci_config else 12,
+        encoder_num_layers=args.devices if args.ci_config else 12,
         hidden_size=768,
         ffn_hidden_size=3072,
         num_attention_heads=12,