FIx SFTTrainer bugs on TRL main (#1276)

* Update sft_trainer.py

* Update trl/trainer/sft_trainer.py

trl/trainer/sft_trainer.py

@@ -254,7 +254,7 @@ def make_inputs_require_grad(module, input, output):
                 data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
         # Pre-process the datasets only once per node. The remaining processes will use the cache.
-        if PartialState().is_local_main_process:
+        with PartialState().local_main_process_first():
             if dataset_kwargs is None:
                 dataset_kwargs = {}
             if train_dataset is not None: