wip

Author: tastelikefeet

src/twinkle/infra/collectors.py

@@ -2,6 +2,7 @@
 from typing import TYPE_CHECKING, Any, Dict, List
 
 from twinkle import DeviceMesh
+from twinkle.utils import pad_and_stack_tensors
 
 if TYPE_CHECKING:
     import torch
@@ -39,7 +40,7 @@ def collect_tensor_dict(outputs: List[Dict[str, Any]], device_mesh: DeviceMesh)
             result[key] = merged
 
         elif isinstance(first_value, torch.Tensor):
-            result[key] = _pad_and_stack_tensors(values)
+            result[key] = pad_and_stack_tensors(values)
 
         elif isinstance(first_value, dict):
             result[key] = collect_tensor_dict(values)
@@ -53,36 +54,3 @@ def collect_tensor_dict(outputs: List[Dict[str, Any]], device_mesh: DeviceMesh)
     if 'loss' in result and len(result['loss']) > 1:
         result['loss'] = np.mean(result['loss'])
     return result
-
-
-def _pad_and_stack_tensors(tensors: List['torch.Tensor'], pad_value: float = -200) -> 'torch.Tensor':
-    import torch
-    if not tensors:
-        raise ValueError('Empty tensor list')
-
-    if len(tensors) == 1:
-        return tensors[0]
-
-    max_ndim = max(t.ndim for t in tensors)
-    expanded_tensors = []
-    for t in tensors:
-        while t.ndim < max_ndim:
-            t = t.unsqueeze(0)
-        expanded_tensors.append(t)
-
-    max_shape = []
-    for dim in range(max_ndim):
-        max_shape.append(max(t.shape[dim] for t in expanded_tensors))
-
-    padded_tensors = []
-    for t in expanded_tensors:
-        if list(t.shape) == max_shape:
-            padded_tensors.append(t)
-        else:
-            pad_params = []
-            for dim in range(max_ndim - 1, -1, -1):
-                pad_params.extend([0, max_shape[dim] - t.shape[dim]])
-            padded = torch.nn.functional.pad(t, pad_params, value=pad_value)
-            padded_tensors.append(padded)
-
-    return torch.cat(padded_tensors, dim=0)

src/twinkle/loss/dpo.py

@@ -284,7 +284,6 @@ def __call__(
         # Extract ref_logps from ref_outputs if provided
         if ref_outputs is not None and ref_logps is None:
             ref_logps = ref_outputs.get('logps')
-
         labels = inputs.get('labels')
         assert labels is not None, "inputs must contain 'labels'"
         if not torch.is_tensor(labels):

src/twinkle/metric/dpo.py

@@ -3,6 +3,7 @@
 from typing import List, Union
 
 from twinkle.data_format import InputFeature, ModelOutput
+from twinkle.utils import pad_and_stack_tensors
 from .base import Metric
 
 
@@ -81,13 +82,20 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M
         """
         import torch
         logps = outputs.get('logps')
-        if logps is None:
+        if logps is None or len(logps) == 0:
             return
+        
+        if isinstance(logps, list) and logps:
+            logps = pad_and_stack_tensors(logps)
 
         # Get labels from inputs
         if isinstance(inputs, list):
-            assert len(inputs) == 1
-            inputs = inputs[0]
+            labels = [input['labels'] for input in inputs]
+            if len(labels) == 1:
+                labels = labels[0]
+            else:
+                labels = pad_and_stack_tensors(labels)
+            inputs = {'labels': labels}
 
         labels = torch.as_tensor(inputs['labels'])
         if labels.dim() == 1:

src/twinkle/model/megatron/megatron.py

@@ -250,6 +250,38 @@ def _not_encoded(inputs):
         assert isinstance(inputs, dict)
         return 'input_ids' not in inputs and 'input_embedding' not in inputs
 
+    @staticmethod
+    def _slice_value_for_microbatch(value, mb_start: int, mb_end: int, micro_batch_size: int):
+        """Recursively slice a value for microbatch processing.
+
+        Handles nested dicts (e.g., ref_outputs: {"logps": tensor}) by recursively
+        slicing internal tensors.
+
+        Args:
+            value: The value to slice (tensor, ndarray, list, dict, or scalar)
+            mb_start: Start index of the microbatch
+            mb_end: End index of the microbatch
+            micro_batch_size: Size of each microbatch
+
+        Returns:
+            Sliced value with the same structure
+        """
+        if isinstance(value, torch.Tensor) and value.dim() >= 1 and value.shape[0] > micro_batch_size:
+            return value[mb_start:mb_end]
+        elif isinstance(value, np.ndarray) and value.ndim >= 1 and value.shape[0] > micro_batch_size:
+            return value[mb_start:mb_end]
+        elif isinstance(value, (list, tuple)) and len(value) > micro_batch_size:
+            return value[mb_start:mb_end]
+        elif isinstance(value, dict):
+            # Recursively slice dict values (e.g., ref_outputs: {"logps": tensor})
+            return {
+                k: MegatronModel._slice_value_for_microbatch(v, mb_start, mb_end, micro_batch_size)
+                for k, v in value.items()
+            }
+        else:
+            # Scalars, small tensors, or non-sliceable values pass through as-is
+            return value
+
     def _postprocess_tensor_cp(self, tensor):
         """All-gather and reconstruct full sequence from CP-split tensor.
 
@@ -401,8 +433,6 @@ def forward_backward(self,
             else:
                 seq_length = original_seq_length
 
-        if 'ref_outputs' in kwargs:
-            breakpoint()
         num_microbatches = len(inputs)
         loss_extra_kwargs_per_mb = []
         if num_microbatches <= 1:
@@ -411,17 +441,10 @@ def forward_backward(self,
             for mb_idx in range(num_microbatches):
                 mb_start = mb_idx * micro_batch_size
                 mb_end = mb_start + micro_batch_size
-                mb_kwargs = {}
-                for key, value in kwargs.items():
-                    if isinstance(value, torch.Tensor) and value.dim() >= 1 and value.shape[0] > micro_batch_size:
-                        mb_kwargs[key] = value[mb_start:mb_end]
-                    elif isinstance(value, np.ndarray) and value.ndim >= 1 and value.shape[0] > micro_batch_size:
-                        mb_kwargs[key] = value[mb_start:mb_end]
-                    elif isinstance(value, (list, tuple)) and len(value) > micro_batch_size:
-                        mb_kwargs[key] = value[mb_start:mb_end]
-                    else:
-                        # Scalars, small tensors, or non-sliceable values pass through as-is
-                        mb_kwargs[key] = value
+                mb_kwargs = {
+                    key: self._slice_value_for_microbatch(value, mb_start, mb_end, micro_batch_size)
+                    for key, value in kwargs.items()
+                }
                 loss_extra_kwargs_per_mb.append(mb_kwargs)
 
         _mb_counter = [0]  # mutable counter for closure

src/twinkle/utils/__init__.py

@@ -10,7 +10,7 @@
 from .parallel import processing_lock
 from .platforms import GPU, NPU, Platform, ensure_hccl_socket_env, ensure_npu_backend
 from .safetensors import LazyTensor, SafetensorLazyLoader, StreamingSafetensorSaver
-from .torch_utils import pad_sequence_to_length, selective_log_softmax, stateless_init_process_group, to_device
+from .torch_utils import pad_sequence_to_length, selective_log_softmax, stateless_init_process_group, to_device, pad_and_stack_tensors
 from .transformers_utils import find_all_linears, find_layers, get_modules_to_not_convert
 from .unsafe import check_unsafe, trust_remote_code
 from .utils import copy_files_by_pattern, deep_getattr

src/twinkle/utils/torch_utils.py

@@ -190,3 +190,36 @@ def stateless_init_process_group(
 
     communicator = Communicator(pg, device=device)
     return communicator
+
+
+def pad_and_stack_tensors(tensors: List['torch.Tensor'], pad_value: float = -200) -> 'torch.Tensor':
+    import torch
+    if not tensors:
+        raise ValueError('Empty tensor list')
+
+    if len(tensors) == 1:
+        return tensors[0]
+
+    max_ndim = max(t.ndim for t in tensors)
+    expanded_tensors = []
+    for t in tensors:
+        while t.ndim < max_ndim:
+            t = t.unsqueeze(0)
+        expanded_tensors.append(t)
+
+    max_shape = []
+    for dim in range(max_ndim):
+        max_shape.append(max(t.shape[dim] for t in expanded_tensors))
+
+    padded_tensors = []
+    for t in expanded_tensors:
+        if list(t.shape) == max_shape:
+            padded_tensors.append(t)
+        else:
+            pad_params = []
+            for dim in range(max_ndim - 1, -1, -1):
+                pad_params.extend([0, max_shape[dim] - t.shape[dim]])
+            padded = torch.nn.functional.pad(t, pad_params, value=pad_value)
+            padded_tensors.append(padded)
+
+    return torch.cat(padded_tensors, dim=0)