diff --git a/open_instruct/dataset_transformation.py b/open_instruct/dataset_transformation.py
index d0ad19f95b..433fb64810 100644
--- a/open_instruct/dataset_transformation.py
+++ b/open_instruct/dataset_transformation.py
@@ -51,7 +51,7 @@
 
 import torch
 import transformers
-from datasets import Dataset, concatenate_datasets, load_dataset
+from datasets import Dataset, concatenate_datasets
 from huggingface_hub import ModelCard, revision_exists
 from rich.console import Console
 from rich.text import Text
@@ -67,6 +67,37 @@
 
 from open_instruct.utils import hf_whoami
 
+# MOVE THIS SOMEWHERE
+def load_dataset(
+    *args, **kwargs
+):
+    from datasets import load_dataset as lds
+    from datasets import DatasetDict
+    import pandas as pd
+
+    try:
+        return lds(*args, **kwargs)
+    except:
+        # - load_dataset sometimes has strict schema
+        # checks and may fail on tools / documents
+        train_files = kwargs['data_files']
+        if isinstance(train_files, str):
+            train_files = [train_files]
+
+        file_type = train_files[0].split('.')[-1]
+        assert file_type in {'json', 'jsonl', 'parquet'}
+
+        dfs = []
+        reader = (
+            partial(pd.read_json , lines=True)
+            if file_type in {'json', 'jsonl'}
+            else partial(pd.read_parquet, engine='auto')
+        )
+        for file in train_files:
+            dfs.append(reader(file, orient='records'))
+
+        raw_datasets = Dataset.from_pandas(pd.concat(dfs))
+        return raw_datasets
 
 # ----------------------------------------------------------------------------
 # Utilities
diff --git a/open_instruct/finetune.py b/open_instruct/finetune.py
index f59a20c8c1..09fce56f59 100644
--- a/open_instruct/finetune.py
+++ b/open_instruct/finetune.py
@@ -235,7 +235,7 @@ class FlatArguments:
             " Only for linear schedulers, currently."
         },
     )
-    weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."})
+    
     weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."})
     timeout: int = field(
         default=1800,
@@ -291,10 +291,7 @@ class FlatArguments:
     load_balancing_weight: float = field(
         default=0.5, metadata={"help": "Weight for load balancing loss if applicable."}
     )
-    clean_checkpoints_at_end: bool = field(
-        default=True, metadata={"help": "Whether to clean up all previous checkpoints at the end of the run."}
-    )
-
+    
     # Experiment tracking
     with_tracking: bool = False
     """If toggled, this experiment will be tracked with Weights and Biases"""
@@ -341,13 +338,7 @@ class FlatArguments:
             "help": "Whether to clean up all previous checkpoints at the end of the run.",
         },
     )
-    final_lr_ratio: Optional[float] = field(
-        default=None,
-        metadata={
-            "help": "Set the final lr value at the end of training to be final_lr_ratio * learning_rate."
-            " Only for linear schedulers, currently."
-        },
-    )
+    
     add_seed_and_date_to_exp_name: bool = True
     additional_model_arguments: Optional[Union[dict, str]] = field(
         default_factory=dict,
diff --git a/open_instruct/model_utils.py b/open_instruct/model_utils.py
index cbeb83cf28..b9d475ed92 100644
--- a/open_instruct/model_utils.py
+++ b/open_instruct/model_utils.py
@@ -394,6 +394,7 @@ def save_with_accelerate(
     output_dir: str,
     use_lora: bool = False,
     model_attribute_to_save: Optional[str] = None,
+    safe_serialization: bool = True,
 ) -> None:
     """`model_attribute_to_save` is for used to save PPO's policy instead of the full model"""
     # set the generation config to an empty setting to be safe.
@@ -429,13 +430,12 @@ def save_with_accelerate(
         if accelerator.is_main_process:
             unwrapped_model.save_pretrained(output_dir, state_dict=state_dict)
     else:
-        # don't use safetensors for saving for now
         unwrapped_model.save_pretrained(
             output_dir,
             is_main_process=accelerator.is_main_process,
             save_function=accelerator.save,
             state_dict=state_dict,
-            safe_serialization=False,
+            safe_serialization=safe_serialization,
         )
 
     if accelerator.is_main_process: