Fix three memory defects in the Megatron path (selective_log_softmax, causal mask, discarded logits)

src/twinkle/model/megatron/megatron.py

@@ -324,11 +324,17 @@ def forward_backward(self,
             # Compatible with DPO
             micro_batch_size = min(2, len(inputs))
         unwrapped_model = self.strategy.unwrap_model(self.model)[0]
+        # No config in the Megatron stack declares ``attention_mask_type``, so this
+        # lookup always fell through to None and the processor could never tell that
+        # the task is causal. Default it for decoder-only causal_lm.
+        attention_mask_type = getattr(unwrapped_model.config, 'attention_mask_type', None)
+        if attention_mask_type is None and task == 'causal_lm':
+            attention_mask_type = 'causal'
         inputs = processor(
             inputs,
             micro_batch_size=micro_batch_size,
             variable_seq_lengths=self.variable_seq_lengths,
-            attention_mask_type=getattr(unwrapped_model.config, 'attention_mask_type', None),
+            attention_mask_type=attention_mask_type,
         )
 
         # Get parallelism settings for sequence padding and splitting
@@ -395,7 +401,11 @@ def post_loss_function(output_tensor, inputs, logps, unpacked_logits=None, entro
                 # 2. PER TOKEN MEAN loss: (gather_sum(per_token_grad * gradient_accumulation_steps))
                 #       / (gradient_accumulation_steps  * world_size ) = avg_per_token_grad
                 counts = torch.tensor(1, device=losses.device)
-            return self.strategy.reduce_loss(losses, counts, output_tensor, logps)
+            # reduce_loss() detaches the logits into the per-microbatch report dict,
+            # which pins every microbatch's full-vocab logits until they are cat'ed
+            # and dropped below. Mirrors the guard in TransformersModel.forward.
+            reported_logits = output_tensor if return_logits else None
+            return self.strategy.reduce_loss(losses, counts, reported_logits, logps)
 
         # Define forward step function for Megatron
         # forward_step_func(data_iterator, model) -> (output_tensor, partial(loss_func))
@@ -523,9 +533,9 @@ def forward_step_func(data_iterator, model):
                 if isinstance(loss_dict, dict):
                     if 'loss' in loss_dict:
                         loss += loss_dict['loss']
-                    if 'logits' in loss_dict:
+                    if loss_dict.get('logits') is not None:
                         logits.append(loss_dict['logits'])
-                    if 'logps' in loss_dict:
+                    if loss_dict.get('logps') is not None:
                         logps.append(loss_dict['logps'])
                     if 'num_tokens' in loss_dict:
                         count += loss_dict['num_tokens']

src/twinkle/processor/base.py

@@ -458,13 +458,19 @@ def prepare_transformers_padding_free_patch(self, inputs: List[InputFeature], **
             _inp['max_length_k'] = int(packed_seq_params.max_seqlen_kv)
         return inputs
 
+    @staticmethod
+    def _backend_builds_causal_mask(attention_mask_type: Optional[str]) -> bool:
+        """True when the accelerator synthesises its own causal mask.
+
+        MindSpeed rebuilds a compressed causal mask for FlashAttention, so
+        handing it the generic dense ``[B, 1, S, S]`` mask is both redundant and
+        expensive at long sequence lengths.
+        """
+        return attention_mask_type == 'causal' and Platform.device_prefix() == 'npu'
+
     def drop_causal_4d_mask(self, inputs: List[InputFeature], **kwargs) -> List[InputFeature]:
-        """On NPU, drop the generic 4D dense mask so MindSpeed can build
-        its own compressed causal mask for FlashAttention."""
-        if Platform.device_prefix() != 'npu':
-            return inputs
-        attention_mask_type = kwargs.get('attention_mask_type')
-        if attention_mask_type != 'causal':
+        """Drop a generic dense mask when the accelerator builds its causal mask."""
+        if not self._backend_builds_causal_mask(kwargs.get('attention_mask_type')):
             return inputs
         for _inp in inputs:
             attention_mask = _inp.get('attention_mask')
@@ -730,7 +736,10 @@ def _fill_optional_sequence_fields(self, batch: List[InputFeature]) -> None:
                         continue
                 feat[key] = torch.full((length, ), pad_value, dtype=torch.long, device=device)
 
-    def _collate_macro_batch(self, inputs: List[InputFeature]) -> InputFeature:
+    def _collate_macro_batch(self,
+                             inputs: List[InputFeature],
+                             *,
+                             omit_megatron_attention_mask: bool = False) -> InputFeature:
         # Work on local copies so squeezing doesn't mutate the caller's original samples.
         squeezed = []
         for _input in inputs:
@@ -782,7 +791,7 @@ def is_mm_position_ids(position_ids):
             for key in text_keys:
                 values = [item[key] for item in text_inputs]
                 if self.framework == 'megatron' and key == 'attention_mask':
-                    result[key] = self._create_4d_attention_mask(values)
+                    result[key] = (None if omit_megatron_attention_mask else self._create_4d_attention_mask(values))
                 elif key == 'position_ids' and is_mm_position_ids(values[0]):
                     result[key] = InputProcessor._pad_sequence(values, self.padding_map[key], self.padding_side)
                     num_axes = values[0].shape[0]
@@ -822,9 +831,11 @@ def collate_fn(self,
                    **kwargs) -> List[InputFeature]:
         if len(inputs) == 1 and self.framework != 'megatron':
             return inputs
+        omit_megatron_attention_mask = (
+            self.framework == 'megatron' and self._backend_builds_causal_mask(kwargs.get('attention_mask_type')))
         if micro_batch_size is None:
             # normal collate
-            outputs = self._collate_macro_batch(inputs)
+            outputs = self._collate_macro_batch(inputs, omit_megatron_attention_mask=omit_megatron_attention_mask)
             for key in outputs:
                 if key in self.VLM_CONCAT_FIELDS:
                     outputs[key] = torch.cat(outputs[key], dim=0)
@@ -835,22 +846,25 @@ def collate_fn(self,
             assert len(inputs) >= micro_batch_size
             outputs = []
             for i in range(0, len(inputs), micro_batch_size):
-                _output = self._collate_macro_batch(inputs[i:i + micro_batch_size])
+                _output = self._collate_macro_batch(
+                    inputs[i:i + micro_batch_size], omit_megatron_attention_mask=omit_megatron_attention_mask)
                 for key in _output:
                     if key in self.VLM_CONCAT_FIELDS:
                         _output[key] = torch.cat(_output[key], dim=0)
                 outputs.append(_output)
             return outputs
         else:
             # each macro batch shares the same length
-            res = self._collate_macro_batch(inputs)
+            res = self._collate_macro_batch(inputs, omit_megatron_attention_mask=omit_megatron_attention_mask)
             keys = list(res.keys())
             outputs = []
             for i in range(0, len(inputs), micro_batch_size):
                 end = i + micro_batch_size
                 output = {}
                 for key in keys:
-                    if key == 'position_ids' and res[key].dim() > 2:
+                    if res[key] is None:
+                        output[key] = None
+                    elif key == 'position_ids' and res[key].dim() > 2:
                         output[key] = res[key][:, i:end, :]
                     elif key in self.VLM_CONCAT_FIELDS:
                         output[key] = torch.cat(res[key][i:end], dim=0)

src/twinkle/utils/torch_utils.py

@@ -59,6 +59,72 @@ def pad_sequence_to_length(
     return F.pad(tensor, pad_tuple, mode='constant', value=pad_value)
 
 
+_CHUNKED_SELECTIVE_LOG_SOFTMAX = None
+
+
+def _chunked_selective_log_softmax_function():
+    """Lazily build the autograd function used by :func:`selective_log_softmax`.
+
+    ``torch`` is imported lazily throughout this module, so the ``Function``
+    subclass cannot live at module scope. It is built once and cached.
+    """
+    global _CHUNKED_SELECTIVE_LOG_SOFTMAX
+    if _CHUNKED_SELECTIVE_LOG_SOFTMAX is not None:
+        return _CHUNKED_SELECTIVE_LOG_SOFTMAX
+    import torch
+
+    class ChunkedSelectiveLogSoftmax(torch.autograd.Function):
+        """One chunk of ``selective_log_softmax`` that recomputes in backward.
+
+        The wide FP32 upcast is a forward-only temporary: only the original
+        dtype chunk and its labels are saved, so no full-vocab FP32 copy of the
+        activation stays resident until the backward pass.
+        """
+
+        @staticmethod
+        def forward(ctx, logits, labels, compute_dtype, return_entropy):
+            chunk = logits.to(compute_dtype)
+            logsumexp = torch.logsumexp(chunk, dim=-1)
+            selected = torch.gather(chunk, dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)
+            logps = selected - logsumexp
+            ctx.save_for_backward(logits, labels)
+            ctx.compute_dtype = compute_dtype
+            ctx.return_entropy = return_entropy
+            if not return_entropy:
+                return logps
+            probabilities = torch.softmax(chunk, dim=-1)
+            return logps, logsumexp - (probabilities * chunk).sum(dim=-1)
+
+        @staticmethod
+        def backward(ctx, *grad_outputs):
+            logits, labels = ctx.saved_tensors
+            compute_dtype = ctx.compute_dtype
+            chunk = logits.to(compute_dtype)
+            probabilities = torch.softmax(chunk, dim=-1)
+            grad_logits = torch.zeros_like(chunk)
+
+            grad_logps = grad_outputs[0]
+            if grad_logps is not None:
+                grad_logps = grad_logps.to(compute_dtype).unsqueeze(-1)
+                # d/dz (z[label] - logsumexp(z)) = onehot(label) - softmax(z)
+                grad_logits.scatter_add_(dim=-1, index=labels.unsqueeze(-1), src=grad_logps)
+                grad_logits.sub_(probabilities * grad_logps)
+
+            if ctx.return_entropy:
+                grad_entropy = grad_outputs[1]
+                if grad_entropy is not None:
+                    grad_entropy = grad_entropy.to(compute_dtype).unsqueeze(-1)
+                    # H(z) = logsumexp(z) - sum(softmax(z) * z)
+                    # dH/dz = softmax(z) * (E_p[z] - z)
+                    expected = (probabilities * chunk).sum(dim=-1, keepdim=True)
+                    grad_logits.add_(probabilities * (grad_entropy * (expected - chunk)))
+
+            return grad_logits.to(logits.dtype), None, None, None
+
+    _CHUNKED_SELECTIVE_LOG_SOFTMAX = ChunkedSelectiveLogSoftmax
+    return _CHUNKED_SELECTIVE_LOG_SOFTMAX
+
+
 def selective_log_softmax(logits, index, return_entropy: bool = False):
     """
     refer: trl/trainer/utils
@@ -85,7 +151,6 @@ def selective_log_softmax(logits, index, return_entropy: bool = False):
         If ``return_entropy`` is True, returns ``(per_token_logps, per_token_entropy)``.
     """
     import torch
-    import torch.nn.functional as F
 
     try:
         from megatron.core import parallel_state as mpu
@@ -100,40 +165,34 @@ def selective_log_softmax(logits, index, return_entropy: bool = False):
     except (ImportError, AssertionError, OSError):
         pass
 
-    if logits.dtype in [torch.float32, torch.float64]:
-        selected_logits = torch.gather(logits, dim=-1, index=index.unsqueeze(-1)).squeeze(-1)
-        if return_entropy:
-            # Per-row loop mirrors the logsumexp path below, to keep peak memory bounded.
-            logsumexp_values = []
-            per_token_entropy = []
-            for row_logits in logits:
-                row_lse = torch.logsumexp(row_logits, dim=-1)
-                logsumexp_values.append(row_lse)
-                # H = lse - E_p[x] = lse - sum(exp(x - lse) * x)
-                row_p = torch.exp(row_logits - row_lse.unsqueeze(-1))
-                per_token_entropy.append(row_lse - (row_p * row_logits).sum(dim=-1))
-            logsumexp_values = torch.stack(logsumexp_values)
-            per_token_entropy = torch.stack(per_token_entropy)
-            per_token_logps = selected_logits - logsumexp_values
-            return per_token_logps, per_token_entropy
-        # loop to reduce peak mem consumption
-        logsumexp_values = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits])
-        per_token_logps = selected_logits - logsumexp_values  # log_softmax(x_i) = x_i - logsumexp(x)
-    else:
-        # logsumexp approach is unstable with bfloat16, fall back to slightly less efficient approach
-        per_token_logps = []
-        per_token_entropy = [] if return_entropy else None
-        for row_logits, row_labels in zip(logits, index, strict=True):  # loop to reduce peak mem consumption
-            row_logps = F.log_softmax(row_logits, dim=-1)
-            row_per_token_logps = row_logps.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze(-1)
-            per_token_logps.append(row_per_token_logps)
-            if return_entropy:
-                # row_logps is already stable; softmax reuses the same numerics.
-                row_p = torch.exp(row_logps)
-                per_token_entropy.append(-(row_p * row_logps).sum(dim=-1))
-        per_token_logps = torch.stack(per_token_logps)
+    vocab_size = logits.shape[-1]
+    flat_logits = logits.reshape(-1, vocab_size)
+    flat_index = index.reshape(-1)
+    compute_dtype = torch.float64 if logits.dtype == torch.float64 else torch.float32
+    element_size = torch.empty((), dtype=compute_dtype).element_size()
+    temporary_bytes = 32 << 20 if return_entropy else 64 << 20
+    chunk_rows = max(1, temporary_bytes // (vocab_size * element_size))
+    chunk_function = _chunked_selective_log_softmax_function()
+
+    logps = []
+    entropies = []
+    for start in range(0, flat_logits.shape[0], chunk_rows):
+        stop = min(start + chunk_rows, flat_logits.shape[0])
+        result = chunk_function.apply(
+            flat_logits[start:stop],
+            flat_index[start:stop],
+            compute_dtype,
+            return_entropy,
+        )
         if return_entropy:
-            return per_token_logps, torch.stack(per_token_entropy)
+            chunk_logps, chunk_entropies = result
+            logps.append(chunk_logps)
+            entropies.append(chunk_entropies)
+        else:
+            logps.append(result)
+    per_token_logps = torch.cat(logps).reshape(index.shape)
+    if return_entropy:
+        return per_token_logps, torch.cat(entropies).reshape(index.shape)
     return per_token_logps
 
 

tests/processor/test_processor.py

@@ -4,6 +4,7 @@
 import torch
 
 import twinkle
+from twinkle import Platform
 from twinkle.processor import InputProcessor
 
 twinkle.initialize(mode='local')
@@ -93,6 +94,48 @@ def test_micro_batch_variable_length(self):
             assert b['input_ids'].shape[0] == 2
 
 
+class TestMegatronCausalMask:
+    """The dense 4D mask must not be built when the backend derives its own."""
+
+    def _collate(self, monkeypatch, device_prefix, attention_mask_type):
+        monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda: device_prefix))
+        proc = InputProcessor(padding_free=False, framework='megatron')
+        proc._create_4d_attention_mask = lambda _: pytest.fail('dense 4D mask allocated')
+        return proc.collate_fn(
+            _make_text_batch(4, seq_len=8),
+            micro_batch_size=2,
+            variable_seq_lengths=False,
+            attention_mask_type=attention_mask_type,
+        )
+
+    def test_causal_npu_omits_dense_mask(self, monkeypatch):
+        outputs = self._collate(monkeypatch, 'npu', 'causal')
+        assert len(outputs) == 2
+        assert all(output['attention_mask'] is None for output in outputs)
+
+    def test_non_causal_still_builds_dense_mask(self, monkeypatch):
+        monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda: 'npu'))
+        proc = InputProcessor(padding_free=False, framework='megatron')
+        outputs = proc.collate_fn(
+            _make_text_batch(4, seq_len=8),
+            micro_batch_size=2,
+            variable_seq_lengths=False,
+            attention_mask_type=None,
+        )
+        assert all(output['attention_mask'].dim() == 4 for output in outputs)
+
+    def test_other_backends_still_build_dense_mask(self, monkeypatch):
+        monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda: 'cuda'))
+        proc = InputProcessor(padding_free=False, framework='megatron')
+        outputs = proc.collate_fn(
+            _make_text_batch(4, seq_len=8),
+            micro_batch_size=2,
+            variable_seq_lengths=False,
+            attention_mask_type='causal',
+        )
+        assert all(output['attention_mask'].dim() == 4 for output in outputs)
+
+
 class TestMultimodalMode:
     """Multimodal: pixel_values, image_grid_thw."""