Make training args fully immutable (#25435)

* Make training args fully immutable

* Working tests, PyTorch

* In test_trainer

* during testing

* Use proper dataclass way

* Fix test

* Another one

* Fix tf

* Lingering slow

* Exception

* Clean

examples/pytorch/image-pretraining/run_mae.py

@@ -163,6 +163,15 @@ class CustomTrainingArguments(TrainingArguments):
         default=1e-3, metadata={"help": "Base learning rate: absolute_lr = base_lr * total_batch_size / 256."}
     )
 
+    def __post_init__(self):
+        # Compute absolute learning rate while args are mutable
+        super().__post_init__()
+        if self.base_learning_rate is not None:
+            total_train_batch_size = self.train_batch_size * self.gradient_accumulation_steps * self.world_size
+            delattr(self, "_frozen")
+            self.learning_rate = self.base_learning_rate * total_train_batch_size / 256
+            setattr(self, "_frozen", True)
+
 
 def collate_fn(examples):
     pixel_values = torch.stack([example["pixel_values"] for example in examples])
@@ -353,13 +362,6 @@ def main():
         # Set the validation transforms
         ds["validation"].set_transform(preprocess_images)
 
-    # Compute absolute learning rate
-    total_train_batch_size = (
-        training_args.train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
-    )
-    if training_args.base_learning_rate is not None:
-        training_args.learning_rate = training_args.base_learning_rate * total_train_batch_size / 256
-
     # Initialize our trainer
     trainer = Trainer(
         model=model,

examples/pytorch/summarization/run_summarization.py

@@ -18,6 +18,7 @@ Fine-tuning the library models for sequence to sequence.
 """
 # You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
 
+import dataclasses
 import logging
 import os
 import sys
@@ -674,14 +675,10 @@ def main():
         return result
 
     # Override the decoding parameters of Seq2SeqTrainer
-    training_args.generation_max_length = (
-        training_args.generation_max_length
-        if training_args.generation_max_length is not None
-        else data_args.val_max_target_length
-    )
-    training_args.generation_num_beams = (
-        data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
-    )
+    if training_args.generation_max_length is None:
+        training_args = dataclasses.replace(training_args, generation_max_length=data_args.val_max_target_length)
+    if training_args.generation_num_beams is None:
+        training_args = dataclasses.replace(training_args, generation_num_beams=data_args.num_beams)
 
     # Initialize our Trainer
     trainer = Seq2SeqTrainer(