fix(train/data): thread dtype through create_collate_fn to create_empty_sample

Address CodeRabbit review feedback on PR #527: the caller of
create_empty_sample — create_collate_fn — had no way to pass a dtype
through, so non-bf16 trainers (e.g. fp32 verifiers, fp16 verifiers) would
still hit a dtype mismatch when the empty-sample fallback fires.

Add a dtype keyword argument to create_collate_fn defaulting to
torch.bfloat16, and thread it into the create_empty_sample call so callers
can match the empty placeholder to their training precision.

The bf16 default keeps existing bf16 training paths unchanged while
unblocking non-bf16 callers cleanly.

Signed-off-by: Thor Lin <coolthor@gmail.com>

Author: coolthor

src/speculators/train/data.py

@@ -434,6 +434,7 @@ def _get_raw_data(self, index):
 def create_collate_fn(
     max_len: int,
     hidden_size: int,
+    dtype: torch.dtype = torch.bfloat16,
     preprocess: Callable[[BatchType], BatchType] | None = None,
 ):
     def collate_fn(batch: list[BatchType | None]) -> BatchType:
@@ -442,8 +443,11 @@ def collate_fn(batch: list[BatchType | None]) -> BatchType:
 
         if not batch:
             # Create empty sample which then gets padded to full
-            # batch size if no valid samples are found
-            batch = [create_empty_sample(hidden_size)]
+            # batch size if no valid samples are found.
+            # Match the configured `dtype` so the placeholder doesn't crash
+            # downstream layers loaded at a different precision (e.g. bf16
+            # weights vs fp32 default placeholders).
+            batch = [create_empty_sample(hidden_size, dtype=dtype)]
 
         collated_data = {}
         for key in batch[0]:  # type: ignore[union-attr]