setup defaults for dataloader to ensure GPU is kept busy (#2632) [skip ci]

src/axolotl/utils/schemas/config.py

@@ -1114,3 +1114,17 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
             raise ValueError("QAT is not supported on torch version < 2.6.0")
 
         return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def default_dataloader_opts(cls, data):
+        if (
+            data.get("dataloader_num_workers") is None
+            and data.get("dataloader_pin_memory") is None
+            and data.get("dataloader_prefetch_factor") is None
+        ):
+            data["dataloader_num_workers"] = data.get("capabilities").get("n_gpu", 1)
+            data["dataloader_pin_memory"] = True
+            data["dataloader_prefetch_factor"] = 256
+
+        return data