Add sample_start_step setting to configure when sampling starts (#949)

Co-authored-by: Jaret Burkett <jaretburkett@gmail.com>

config/examples/modal/modal_train_lora_flux_24gb.yaml

@@ -70,6 +70,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/modal/modal_train_lora_flux_schnell_24gb.yaml

@@ -72,6 +72,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_flex_redux.yaml

@@ -85,6 +85,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         # I leave half blank to test prompt and unprompted

config/examples/train_full_fine_tune_flex.yaml

@@ -81,6 +81,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_full_fine_tune_lumina.yaml

@@ -73,6 +73,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_chroma_24gb.yaml

@@ -78,6 +78,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flex2_24gb.yaml

@@ -129,6 +129,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flex_24gb.yaml

@@ -75,6 +75,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flux_24gb.yaml

@@ -70,6 +70,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flux_kontext_24gb.yaml

@@ -79,6 +79,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_flux_schnell_24gb.yaml

@@ -72,6 +72,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_hidream_48.yaml

@@ -86,6 +86,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_lumina.yaml

@@ -70,6 +70,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_omnigen2_24gb.yaml

@@ -68,6 +68,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_qwen_image_24gb.yaml

@@ -69,6 +69,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_qwen_image_edit_2509_32gb.yaml

@@ -70,6 +70,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         # you can provide up to 3 control images here

config/examples/train_lora_qwen_image_edit_32gb.yaml

@@ -68,6 +68,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         samples:

config/examples/train_lora_sd35_large_24gb.yaml

@@ -71,6 +71,7 @@ config:
       sample:
         sampler: "flowmatch" # must match train.noise_scheduler
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         prompts:

config/examples/train_lora_wan21_14b_24gb.yaml

@@ -80,6 +80,7 @@ config:
       sample:
         sampler: "flowmatch"
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 832
         height: 480
         num_frames: 40

config/examples/train_lora_wan21_1b_24gb.yaml

@@ -69,6 +69,7 @@ config:
       sample:
         sampler: "flowmatch"
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 832
         height: 480
         num_frames: 40

config/examples/train_lora_wan22_14b_24gb.yaml

@@ -80,6 +80,7 @@ config:
       sample:
         sampler: "flowmatch"
         sample_every: 250 # sample every this many steps
+        sample_start_step: 0 # start sampling at this step
         width: 1024
         height: 1024
         # set to 1 for images