Just to see the diff

examples/research_projects/codeparrot/requirements.txt

@@ -1,9 +1,9 @@
-transformers==4.19.0
+git+https://github.com/loubnabnl/transformers.git@loss-reduction-none
+accelerate==0.15.0
 datasets==1.16.0
 wandb==0.12.0
 tensorboard==2.6.0
 torch==1.11.0
 huggingface-hub==0.1.0
-git+https://github.com/huggingface/accelerate.git@3c45b6f760ad8745be9ebc9bbb26f5b04dea4abe
 datasketch==1.5.7
 dpu_utils

examples/research_projects/codeparrot/scripts/arguments.py

@@ -20,12 +20,16 @@ class TrainingArguments:
     dataset_name_valid: Optional[str] = field(
         default="codeparrot/codeparrot-clean-valid", metadata={"help": "Name or path of validation dataset."}
     )
-    train_batch_size: Optional[int] = field(default=2, metadata={"help": "Batch size for training."})
-    valid_batch_size: Optional[int] = field(default=2, metadata={"help": "Batch size for evaluation."})
+    train_batch_size: Optional[int] = field(default=320, metadata={"help": "Batch size for training."})
+    train_batch_size_select: Optional[int] = field(default=32, metadata={"help": "Batch size to subselect for training."})
+    valid_batch_size: Optional[int] = field(default=32, metadata={"help": "Batch size for evaluation."})
     weight_decay: Optional[float] = field(default=0.1, metadata={"help": "Value of weight decay."})
     shuffle_buffer: Optional[int] = field(
         default=10000, metadata={"help": "Size of buffer used to shuffle streaming dataset."}
     )
+    no_streaming: Optional[bool] = field(
+        default=False, metadata={"help": "Whether not to use streaming for the dataset."}
+    )
     learning_rate: Optional[float] = field(default=2e-4, metadata={"help": "Learning rate fo training."})
     lr_scheduler_type: Optional[str] = field(default="cosine", metadata={"help": "Learning rate."})
     num_warmup_steps: Optional[int] = field(
@@ -47,6 +51,18 @@ class TrainingArguments:
         default=1024,
         metadata={"help": "Interval to save checkpoints. Measured as number of forward passes not training steps."},
     )
+    selection_method: Optional[str] = field(
+        default=None, metadata={"help": "Selection method to subselect from the batch size. Can be uniform or rholoss"}
+    )
+    irred_losses: Optional[str] = field(
+        default="irred_losses.pt", metadata={"help": "Path to irreducible losses pt file. Must be supplied if selection_method is rholoss"}
+    )
+    compute_irred_losses: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "If True irreducible losses are computed and saved to the path specified by irred_losses."
+        },
+    )
     resume_from_checkpoint: Optional[str] = field(
         default=None, metadata={"help": "States path if the training should continue from a checkpoint folder."}
     )