fix bugs

src/twinkle/infra/__init__.py

@@ -546,7 +546,11 @@ def _new_init_body(self, _caller, *args, **kwargs):
                         kwargs['min_batch_size'] = device_mesh.data_world_size
                     init_method(self, *args, **kwargs)
                 else:
-                    # Pop the device_mesh
+                    if device_mesh is not None:
+                        logger.warning(f'{cls.__name__} was given a device_mesh but it is being DROPPED: twinkle '
+                                       'holds no device group, so call twinkle.initialize(...) before '
+                                       'constructing it. Training will otherwise run with device_mesh=None '
+                                       '(single-rank loss normalisation and no metric aggregation).')
                     args = [arg for arg in args if not isinstance(arg, DeviceMesh)]
                     kwargs = {key: value for key, value in kwargs.items() if not isinstance(value, DeviceMesh)}
                     init_method(self, *args, **kwargs)

src/twinkle/model/transformers/transformers.py

@@ -947,6 +947,10 @@ def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str, Optimizer], *
                 Any parameters needed to construct the optimizer instance.
         """
         adapter_name = kwargs.pop('adapter_name', self._get_default_group())
+        # Metric copies the dp group at construction, and OptimizerGroup builds metrics before
+        # dist init -- so rebuild here (first path that runs post-init on every backend) to get
+        # dp-wide token-weighted logging. Logging only; gradients are unaffected.
+        self._ensure_optimizer_dp_groups()
         optimizer_config = self.optimizer_group[adapter_name]
         if isinstance(optimizer_cls, Optimizer):
             optimizer_config.optimizer = optimizer_cls

src/twinkle/processor/base.py

@@ -102,7 +102,6 @@ def prepare_outputs(self, inputs: List[InputFeature], **kwargs) -> Union[List[In
     def prepare_inputs(self, inputs: Union[List[InputFeature], InputFeature], **kwargs) -> List[InputFeature]:
 
         def to_tensor(_input):
-            import torch
             for key in list(_input.keys()):
                 value = _input[key]
                 # Ray/pyarrow can return numpy or list scalars; normalize to tensors.
@@ -473,7 +472,6 @@ def _pad_sequence(sequences, padding_value, padding_side, concat=None):
                 sequences, pad_value=padding_value, concat=concat if concat is not None else (sequences[0].dim() >= 2))
         else:
             # left padding
-            import torch
             max_len = max([s.shape[0] for s in sequences])
 
             padded_sequences = []
@@ -486,7 +484,6 @@ def _pad_sequence(sequences, padding_value, padding_side, concat=None):
 
     @staticmethod
     def _create_4d_attention_mask(attention_mask):
-        import torch
         seq_lens = [s.shape[0] for s in attention_mask]
         max_len = max(seq_lens)
         device = attention_mask[0].device
@@ -652,7 +649,6 @@ def unpack_inputs(self, inputs: List[Dict[str, Any]], task: str = 'causal_lm') -
 
     @staticmethod
     def to_transformers_dict(inputs: List[InputFeature], **kwargs) -> List[InputFeature]:
-        import torch
         results = []
         for _input in inputs:
             output = {}
@@ -671,6 +667,8 @@ def to_transformers_dict(inputs: List[InputFeature], **kwargs) -> List[InputFeat
                 'max_length_k',
                 'packed_seq_params',
                 'routed_experts',
+                'mm_token_type_ids',
+                'second_per_grid_ts',
             ] + list(InputProcessor.VLM_CONCAT_FIELDS)
             for key in list(_input.keys()):
                 if key not in _keys:
@@ -683,13 +681,43 @@ def to_transformers_dict(inputs: List[InputFeature], **kwargs) -> List[InputFeat
             results.append(InputFeature(**output))
         return results
 
-    def _collate_macro_batch(self, inputs: List[InputFeature]) -> InputFeature:
-        import torch
+    # when training on mixed text + multimodal data, some fields (e.g. mm_token_type_ids)
+    # only exist on multimodal samples. We pad those fields for the missing samples instead of
+    # letting collate blow up with a KeyError.
+    def _fill_optional_sequence_fields(self, batch: List[InputFeature]) -> None:
+        if len(batch) < 2:
+            return
+        seq_fields = set(self.padding_map) - set(self.VLM_CONCAT_FIELDS)
+        present = {k for feat in batch for k in feat if k in seq_fields}
+        for key in present:
+            missing = [feat for feat in batch if feat.get(key) is None]
+            if not missing or len(missing) == len(batch):
+                continue  # all-present (no gap) or all-absent (nothing to align to) -> leave as is
+            pad_value = self.padding_map[key]
+            for feat in missing:
+                input_ids = feat.get('input_ids')
+                if input_ids is None:
+                    continue
+                # Match the row's device/dtype: prepare_inputs may already have moved rows to the
+                # accelerator, and a CPU fill would break the cat inside the collate below.
+                reference = input_ids if isinstance(input_ids, torch.Tensor) else None
+                length = reference.shape[-1] if reference is not None else len(input_ids)
+                feat[key] = torch.full((length, ),
+                                       pad_value,
+                                       dtype=torch.long,
+                                       device=reference.device if reference is not None else None)
 
+    def _collate_macro_batch(self, inputs: List[InputFeature]) -> InputFeature:
+        # Work on local copies so squeezing doesn't mutate the caller's original samples.
+        squeezed = []
         for _input in inputs:
+            _input = dict(_input)
             for key in list(_input.keys()):
                 if isinstance(_input[key], torch.Tensor):
                     _input[key] = _input[key].squeeze()
+            squeezed.append(_input)
+        inputs = squeezed
+        self._fill_optional_sequence_fields(inputs)
 
         vlm_fields = {k: [] for k in self.VLM_CONCAT_FIELDS}
         text_inputs = []