Add in validation for compile

Author: joecummings

apps/sft/main.py

@@ -95,6 +95,13 @@ def record_batch_metrics(self, data_metrics: list):
 
     @endpoint
     async def setup(self):
+        # Validate that compile is only used with flex attention
+        if self.job_config.training.compile:
+            raise ValueError(
+                "training.compile=True is not currently supported. "
+                "Compile is only supported with flex attention enabled, which requires PyTorch nightly. "
+                "Please set training.compile=false in your config."
+            )
 
         # all ranks should record loss, except when PP=True. Then, only the last stage should record loss.
         self.rank_should_record_loss = True
@@ -150,6 +157,7 @@ def setup_data(self, dataset_configs: list[dict]) -> StatefulDataLoader:
         Raises:
             ValueError: If multiple datasets provided (not yet supported)
         """
+
         # TODO felipemello: Currently only support single dataset
         if len(dataset_configs) > 1:
             raise ValueError(

src/forge/data/collate.py

@@ -46,7 +46,9 @@ def collate_padded(batch: list[dict[str, Any]]) -> dict[str, Any]:
         tokens_list.append(padded_tokens)
 
         # Pad labels with CROSS_ENTROPY_IGNORE_IDX (-100)
-        padded_labels = F.pad(sample["labels"], (0, pad_len), value=CROSS_ENTROPY_IGNORE_IDX)
+        padded_labels = F.pad(
+            sample["labels"], (0, pad_len), value=CROSS_ENTROPY_IGNORE_IDX
+        )
         labels_list.append(padded_labels)
 
     # Stack into batch

tests/unit_tests/datasets/test_packed.py

@@ -1125,4 +1125,3 @@ def test_metrics_flattened(self):
         assert "metrics" in result
         # Should be flattened from [[metric1, metric2], [metric3]] to [metric1, metric2, metric3]
         assert len(result["metrics"]) == 3
-