Reduce SFT fine-tuning steps to avoid timeout (GoogleCloudPlatform#831)

SurbhiJainUSC · web-flow · commit e155edaade9a · 2025-08-18T12:10:22.000-07:00
diff --git a/dags/multipod/maxtext_sft_trainer.py b/dags/multipod/maxtext_sft_trainer.py
@@ -47,7 +47,7 @@
         'export PRE_TRAINED_MODEL_TOKENIZER=meta-llama/Llama-2-7b-chat-hf',
         'export PRE_TRAINED_MODEL_CKPT_PATH=gs://maxtext-model-checkpoints/llama2-7b-chat/scanned/0/items',
         f'export BASE_OUTPUT_DIRECTORY={base_output_directory}',
-        'export STEPS=2500',
+        'export STEPS=1000',
         'export PROMPT="Suggest some famous landmarks in London."',
         'bash end_to_end/tpu/test_sft_trainer.sh',
     )

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@`
`47`	`47`	`'export PRE_TRAINED_MODEL_TOKENIZER=meta-llama/Llama-2-7b-chat-hf',`
`48`	`48`	`'export PRE_TRAINED_MODEL_CKPT_PATH=gs://maxtext-model-checkpoints/llama2-7b-chat/scanned/0/items',`
`49`	`49`	`f'export BASE_OUTPUT_DIRECTORY={base_output_directory}',`
`50`		`- 'export STEPS=2500',`
	`50`	`+ 'export STEPS=1000',`
`51`	`51`	`'export PROMPT="Suggest some famous landmarks in London."',`
`52`	`52`	`'bash end_to_end/tpu/test_sft_trainer.sh',`
`53`	`53`	`)`