diff --git a/src/twinkle/infra/__init__.py b/src/twinkle/infra/__init__.py index a2760c90..a3c90eae 100644 --- a/src/twinkle/infra/__init__.py +++ b/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) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index ffbcebcc..57fc0ee6 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/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 diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 483fb821..2725b99d 100644 --- a/src/twinkle/processor/base.py +++ b/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 = []