fix tests x2

Author: eshwarprasadS

src/instructlab/training/batch_loss_manager.py

@@ -77,10 +77,12 @@ def process_batch(self, batch: list[CollatedItem]) -> tuple[BatchMetrics, float]
         num_minibatches = len(batch)
 
         # initialize accumulation variables
-        batch_total_samples = 0.0
-        batch_total_length = 0.0
-        accumulated_loss = 0.0
-        accumulated_aux_loss = 0.0 if self.model.is_gpt_oss else None
+        batch_total_samples = 0
+        batch_total_length = 0
+        accumulated_loss = torch.tensor([0.0], dtype=torch.float32)
+        accumulated_aux_loss = (
+            torch.tensor([0.0], dtype=torch.float32) if self.model.is_gpt_oss else None
+        )
         grad_accum_steps = 0
 
         # process each minibatch
@@ -134,22 +136,25 @@ def process_batch(self, batch: list[CollatedItem]) -> tuple[BatchMetrics, float]
     def _prepare_model_inputs(self, mb: CollatedItem) -> ModelInputs:
         """Prepare and move model inputs to GPU."""
         model_inputs = ModelInputs(
-            input_ids=mb["input_ids"],
-            labels=mb["labels"],
-            position_ids=mb["position_ids"],
+            input_ids=mb["input_ids"].to(device=self.torch_device),
+            labels=mb["labels"].to(device=self.torch_device),
         )
-        if "attention_mask" in mb:
-            model_inputs["attention_mask"] = mb["attention_mask"]
 
-        # send tensors to gpu
-        for k in model_inputs.keys():
-            model_inputs[k] = model_inputs[k].to(device=self.torch_device)
+        # add optional fields onto `model_inputs` object
+        if "attention_mask" in mb:
+            model_inputs["attention_mask"] = mb["attention_mask"].to(
+                device=self.torch_device
+            )
+        if "position_ids" in mb:
+            model_inputs["position_ids"] = mb["position_ids"].to(
+                device=self.torch_device
+            )
 
         return model_inputs
 
     def _reduce_metrics(
-        self, batch_total_samples: float, batch_total_length: float
-    ) -> tuple[float, float]:
+        self, batch_total_samples: int, batch_total_length: int
+    ) -> tuple[int, int]:
         """Reduce rank-specific metrics across devices."""
         inputs_to_reduce = torch.tensor(
             [batch_total_samples, batch_total_length],

src/instructlab/training/sampler.py

@@ -249,8 +249,10 @@ def __init__(
         else:
             self.batch_packer = batch_lengths_to_minibatches_padded
             # Create a wrapper for padded collate that includes pad_token_id
-            self.collate_fn = lambda mb, tokens: padded_mb_collate_fn(
-                mb, tokens, pad_token_id
+            self.collate_fn = (
+                lambda minibatch, batch_num_loss_counted_tokens: padded_mb_collate_fn(
+                    minibatch, batch_num_loss_counted_tokens, pad_token_id
+                )
             )
 
     def __call__(self, batch: list[dict]):

src/instructlab/training/type_definitions.py

@@ -68,7 +68,7 @@ class ModelInputs(t.TypedDict):
 
     input_ids: Required[torch.Tensor]
     labels: Required[torch.Tensor]
-    position_ids: Required[torch.Tensor]
+    position_ids: NotRequired[torch.Tensor]
     attention_mask: NotRequired[torch.Tensor]  # used when not training in padding free
 
 

tests/unit/test_data_process.py

@@ -890,9 +890,9 @@ def test_with_qwen_tokenizer(self):
         self.assertNotIn(unmask_begin_id, result["input_ids"])
         self.assertNotIn(unmask_end_id, result["input_ids"])
 
-    def test_with_mistral_tokenizer(self):
-        """Test reasoning_content functionality with Mistral tokenizer."""
-        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
+    def test_with_phi_tokenizer(self):
+        """Test reasoning_content functionality with Phi-4 tokenizer."""
+        tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-4-mini-instruct")
 
         # Add the unmask tokens to the tokenizer
         tokenizer.add_special_tokens(