Add force sample toggle to the ui

ui/src/app/jobs/new/SimpleJob.tsx

@@ -809,7 +809,28 @@ export default function SimpleJob({
                       label="Skip First Sample"
                       className="pt-4"
                       checked={jobConfig.config.process[0].train.skip_first_sample || false}
-                      onChange={value => setJobConfig(value, 'config.process[0].train.skip_first_sample')}
+                      onChange={value => {
+                        setJobConfig(value, 'config.process[0].train.skip_first_sample');
+                        // cannot do both, so disable the other
+                        if (value){
+                          setJobConfig(false, 'config.process[0].train.force_first_sample');
+                        }
+                      }}
+                    />
+                  </div>
+                  <div>
+                    <Checkbox
+                      label="Force First Sample"
+                      className="pt-1"
+                      checked={jobConfig.config.process[0].train.force_first_sample || false}
+                      docKey={'train.force_first_sample'}
+                      onChange={value => {
+                        setJobConfig(value, 'config.process[0].train.force_first_sample');
+                        // cannot do both, so disable the other
+                        if (value){
+                          setJobConfig(false, 'config.process[0].train.skip_first_sample');
+                        }
+                      }}
                     />
                   </div>
                   <div>
@@ -817,7 +838,13 @@ export default function SimpleJob({
                       label="Disable Sampling"
                       className="pt-1"
                       checked={jobConfig.config.process[0].train.disable_sampling || false}
-                      onChange={value => setJobConfig(value, 'config.process[0].train.disable_sampling')}
+                      onChange={value => {
+                        setJobConfig(value, 'config.process[0].train.disable_sampling');
+                        // cannot do both, so disable the other
+                        if (value){
+                          setJobConfig(false, 'config.process[0].train.force_first_sample');
+                        }
+                      }}
                     />
                   </div>
                 </FormGroup>

ui/src/app/jobs/new/jobConfig.ts

@@ -75,6 +75,7 @@ export const defaultJobConfig: JobConfig = {
             ema_decay: 0.99,
           },
           skip_first_sample: false,
+          force_first_sample: false,
           disable_sampling: false,
           dtype: 'bf16',
           diff_output_preservation: false,

ui/src/docs.tsx

@@ -155,6 +155,17 @@ const docs: { [key: string]: ConfigDoc } = {
       </>
     ),
   },
+  'train.force_first_sample': {
+    title: 'Force First Sample',
+    description: (
+      <>
+        This option will force the trainer to generate samples when it starts. The trainer will normally only generate a first sample
+        when nothing has been trained yet, but will not do a first sample when resuming from an existing checkpoint. This option
+        forces a first sample every time the trainer is started. This can be useful if you have changed sample prompts and want to see
+        the new prompts right away.
+      </>
+    ),
+  },
 };
 
 export const getDoc = (key: string | null | undefined): ConfigDoc | null => {

ui/src/types.ts

@@ -117,6 +117,7 @@ export interface TrainConfig {
     weight_decay: number;
   };
   skip_first_sample: boolean;
+  force_first_sample: boolean;
   disable_sampling: boolean;
   diff_output_preservation: boolean;
   diff_output_preservation_multiplier: number;