Record training loss (#539)

Per the PR title, this PR solves issue
#488.

modyn/common/grpc/grpc_helpers.py

@@ -216,40 +216,38 @@ def prepare_start_training_request(
 
         grad_scaler_config = training_config.grad_scaler_config if training_config.grad_scaler_config else {}
 
-        start_training_kwargs = {
-            "pipeline_id": pipeline_id,
-            "trigger_id": trigger_id,
-            "device": training_config.device,
-            "use_pretrained_model": previous_model_id is not None,
-            "pretrained_model_id": previous_model_id or -1,
-            "load_optimizer_state": False,  # TODO(#137): Think about this.
-            "batch_size": training_config.batch_size,
-            "torch_optimizers_configuration": TrainerServerJsonString(value=json.dumps(optimizers_config)),
-            "torch_criterion": training_config.optimization_criterion.name,
-            "criterion_parameters": TrainerServerJsonString(value=criterion_config),
-            "data_info": Data(
+        return StartTrainingRequest(
+            pipeline_id=pipeline_id,
+            trigger_id=trigger_id,
+            device=training_config.device,
+            use_pretrained_model=previous_model_id is not None,
+            pretrained_model_id=previous_model_id or -1,
+            load_optimizer_state=False,  # TODO(#137): Think about this.
+            batch_size=training_config.batch_size,
+            torch_optimizers_configuration=TrainerServerJsonString(value=json.dumps(optimizers_config)),
+            torch_criterion=training_config.optimization_criterion.name,
+            criterion_parameters=TrainerServerJsonString(value=criterion_config),
+            data_info=Data(
                 dataset_id=data_config.dataset_id,
                 num_dataloaders=training_config.dataloader_workers,
             ),
-            "checkpoint_info": checkpoint_info,
-            "transform_list": data_config.transformations,
-            "bytes_parser": PythonString(value=data_config.bytes_parser_function),
-            "label_transformer": PythonString(value=data_config.label_transformer_function),
-            "lr_scheduler": TrainerServerJsonString(value=json.dumps(lr_scheduler_configs)),
-            "grad_scaler_configuration": TrainerServerJsonString(value=json.dumps(grad_scaler_config)),
-            "epochs_per_trigger": training_config.epochs_per_trigger,
-            "num_prefetched_partitions": training_config.num_prefetched_partitions,
-            "parallel_prefetch_requests": training_config.parallel_prefetch_requests,
-            "seed": training_config.seed,
-            "tokenizer": PythonString(value=tokenizer) if tokenizer is not None else None,
-            "num_samples_to_pass": num_samples_to_pass,
-            "shuffle": training_config.shuffle,
-            "enable_accurate_gpu_measurements": training_config.enable_accurate_gpu_measurements,
-        }
-
-        cleaned_kwargs: dict[str, Any] = {k: v for k, v in start_training_kwargs.items() if v is not None}
-
-        return StartTrainingRequest(**cleaned_kwargs)
+            checkpoint_info=checkpoint_info,
+            transform_list=data_config.transformations,
+            bytes_parser=PythonString(value=data_config.bytes_parser_function),
+            label_transformer=PythonString(value=data_config.label_transformer_function),
+            lr_scheduler=TrainerServerJsonString(value=json.dumps(lr_scheduler_configs)),
+            grad_scaler_configuration=TrainerServerJsonString(value=json.dumps(grad_scaler_config)),
+            epochs_per_trigger=training_config.epochs_per_trigger,
+            num_prefetched_partitions=training_config.num_prefetched_partitions,
+            parallel_prefetch_requests=training_config.parallel_prefetch_requests,
+            seed=training_config.seed,  # seed is an optional field which can accept None
+            # tokenizer is an optional field which can accept None
+            tokenizer=PythonString(value=tokenizer) if tokenizer is not None else None,
+            num_samples_to_pass=num_samples_to_pass if num_samples_to_pass is not None else 0,
+            shuffle=training_config.shuffle,
+            enable_accurate_gpu_measurements=training_config.enable_accurate_gpu_measurements,
+            record_loss_every=training_config.record_loss_every,
+        )
 
     def start_training(
         self,

modyn/config/schema/pipeline/training/config.py

@@ -131,6 +131,10 @@ class TrainingConfig(ModynBaseModel):
         description="The ID of the model that should be used as the initial model.",
     )
     checkpointing: CheckpointingConfig = Field(description="Configuration of checkpointing during training")
+    record_loss_every: int = Field(
+        default=0,
+        description="Record the training loss in the trainer_log very n-th batch/step. If 0, loss is not recorded.",
+    )
     optimizers: List[OptimizerConfig] = Field(
         description="An array of the optimizers for the training",
         min_length=1,

modyn/protos/trainer_server.proto

@@ -57,6 +57,7 @@ message StartTrainingRequest {
   int64 num_samples_to_pass = 23;
   bool shuffle = 24;
   bool enable_accurate_gpu_measurements = 25;
+  int64 record_loss_every = 26;
 }
 
 message StartTrainingResponse {

modyn/tests/common/grpc/test_grpc_helpers.py

@@ -89,6 +89,8 @@ def test_prepare_start_training_request(
     # for bool value False is the default value so we don't need to test it
     pipeline_training_config.shuffle = True
     pipeline_training_config.enable_accurate_gpu_measurements = True
+    # for int value 0 is the default value so we don't need to test it
+    pipeline_training_config.record_loss_every = 10
     pipeline_training_config.optimization_criterion.config = {"key": "value"}
     pipeline_training_config.use_previous_model = previous_model_id is not None
 
@@ -136,6 +138,7 @@ def test_prepare_start_training_request(
     assert req.num_samples_to_pass == (num_samples_to_pass if num_samples_to_pass is not None else 0)
     assert req.shuffle
     assert req.enable_accurate_gpu_measurements
+    assert req.record_loss_every == 10
 
 
 @patch("modyn.common.grpc.grpc_helpers.grpc_connection_established", return_value=True)

modyn/tests/trainer_server/internal/trainer/test_pytorch_trainer.py

@@ -737,6 +737,8 @@ def test_train(
     status = status_queue.get()
     assert status["num_batches"] == 0
     assert status["num_samples"] == 0
+    # we didn't enable recording the training loss
+    assert len(trainer._log["training_loss"]) == 0
     status_state = torch.load(io.BytesIO(status_queue.get()))
     checkpointed_state = {
         "model": OrderedDict(