From 1797188d171da4c987377f594175f4712ff2bff8 Mon Sep 17 00:00:00 2001 From: Wings-Of-Disaster <1327040545@qq.com> Date: Mon, 14 Sep 2026 16:41:13 +0800 Subject: [PATCH 1/3] fix(utils): recompute selective_log_softmax in backward instead of pinning it The per-row loop was documented as 'loop to reduce peak mem consumption', but under autograd every row's full-vocab intermediate is saved for backward, so the loop pins one copy of the vocabulary per row (two when return_entropy is set) until the backward pass finishes. At 16k tokens and a 151936-word vocab that is 4.6 GiB of extra resident activation per micro-batch in bfloat16, plus tens of thousands of tiny allocations and kernel launches. Replace the loop with a chunked autograd.Function that computes in float32 and saves only the original-dtype chunk and its labels, recomputing the softmax in backward. Nothing wider than the labels is retained, and the float32 compute also removes the bfloat16 quantization of the log-probabilities (~0.2 ULP at the magnitudes a peaked row produces). --- src/twinkle/utils/torch_utils.py | 127 ++++++++++++++++++++++--------- tests/utils/test_utils.py | 62 +++++++++++++++ 2 files changed, 155 insertions(+), 34 deletions(-) diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 88a787ec..513a3a36 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/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 diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index 556b8dc8..cff25b4e 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -144,6 +144,68 @@ def test_bfloat16_fallback(self): # bfloat16 resolution rather than float32's. assert torch.allclose(result.float(), expected, atol=5e-2) + def test_bfloat16_does_not_quantize_large_negative_logps(self): + # A peaked row pushes the selected log-prob to ~-100.3, where the + # bfloat16 ULP is 0.5. Computing the softmax in float32 and only then + # selecting keeps the result at float32 resolution. + vocab_size = 4096 + logits = torch.zeros(2, vocab_size, dtype=torch.bfloat16) + logits[:, 0] = 100.0 + logits[:, 1] = -0.3 + index = torch.ones(2, dtype=torch.long) + + result = selective_log_softmax(logits, index) + expected = torch.gather( + logits.double().log_softmax(-1), -1, index.unsqueeze(-1)).squeeze(-1) + + assert torch.allclose(result.double(), expected, atol=1e-3) + + @pytest.mark.parametrize('dtype', [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize('return_entropy', [False, True]) + def test_backward_saves_no_extra_full_vocab_activation(self, dtype, return_entropy): + # Every tensor handed to save_for_backward stays resident from the end of + # forward until backward runs, so a per-row copy of the vocabulary is as + # expensive as materializing the whole log_softmax. Recomputing in + # backward must keep nothing wider than the labels. + rows, vocab_size = 8, 512 + logits = torch.randn(rows, vocab_size, dtype=dtype, requires_grad=True) + index = torch.randint(0, vocab_size, (rows, )) + logits_storage = logits.untyped_storage().data_ptr() + extra_storages = {} + + def pack(tensor): + storage = tensor.untyped_storage() + if storage.data_ptr() != logits_storage: + extra_storages[storage.data_ptr()] = storage.nbytes() + return tensor + + with torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor): + result = selective_log_softmax(logits, index, return_entropy=return_entropy) + + outputs = result if return_entropy else (result, ) + torch.autograd.backward(outputs, [torch.ones_like(output) for output in outputs]) + + vocab_bytes = vocab_size * logits.element_size() + assert all(size < vocab_bytes for size in extra_storages.values()) + + def test_scales_with_chunks_not_with_rows(self): + # The per-row Python loop issued one save per row per intermediate, i.e. + # thousands of tiny kernel launches at production sequence lengths. + rows, vocab_size = 512, 512 + logits = torch.randn(rows, vocab_size, requires_grad=True) + index = torch.randint(0, vocab_size, (rows, )) + save_calls = [] + + def pack(tensor): + save_calls.append(tensor) + return tensor + + with torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor): + result = selective_log_softmax(logits, index) + result.sum().backward() + + assert len(save_calls) <= 32 + class TestPadAndStackTensors: From a46df5510df89d1c03351248e3dfe62fcd406bc7 Mon Sep 17 00:00:00 2001 From: Wings-Of-Disaster <1327040545@qq.com> Date: Mon, 14 Sep 2026 16:41:19 +0800 Subject: [PATCH 2/3] fix(processor): stop building the dense 4D causal mask the backend rebuilds Two defects combined to make drop_causal_4d_mask dead code: - MegatronModel.forward_backward read attention_mask_type off the model config, but neither mcore TransformerConfig nor mcore_bridge ModelConfig declares that attribute, so the processor always received None and its 'attention_mask_type != causal' guard always returned early. - Even when it did fire, the hook runs four pipeline steps after collate_fn, so the [B, 1, S, S] bool mask had already been allocated (256 MiB per micro-batch at S=16384, roughly twice that in peak with the tril/invert temporaries) only to be set to None. Default the type to 'causal' for decoder-only causal_lm and decide in collate_fn itself, so _create_4d_attention_mask is never called on the NPU/MindSpeed path. The hook stays as a fallback for callers that build the mask elsewhere. --- src/twinkle/model/megatron/megatron.py | 8 ++++- src/twinkle/processor/base.py | 38 ++++++++++++++++------- tests/processor/test_processor.py | 43 ++++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 13 deletions(-) diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 851529c6..0b722b57 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/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 diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 131ae81c..04a7bb21 100644 --- a/src/twinkle/processor/base.py +++ b/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,7 +846,8 @@ 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) @@ -843,14 +855,16 @@ def collate_fn(self, 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) diff --git a/tests/processor/test_processor.py b/tests/processor/test_processor.py index d730499d..a2be36ea 100644 --- a/tests/processor/test_processor.py +++ b/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.""" From dc924457ab0ad70333c18f339b5420d3fb00d9d5 Mon Sep 17 00:00:00 2001 From: Wings-Of-Disaster <1327040545@qq.com> Date: Mon, 14 Sep 2026 16:41:23 +0800 Subject: [PATCH 3/3] fix(megatron): do not retain and cat logits that are discarded reduce_loss() unconditionally puts logits.detach() into the per-microbatch report dict, forward_backward appends every one of them, concatenates the whole set into a second full-size buffer, and then drops it one line later because return_logits defaults to False. At 16k tokens per micro-batch, two micro-batches and a 151936-word vocab that is an 18.6 GiB peak for bfloat16 logits nothing ever reads, held across the backward pass where memory is tightest. Report the logits only when the caller asked for them, mirroring the guard TransformersModel.forward has had since 6921342, and skip None entries when collecting the report dicts so the same path cannot cat a None. The 'logps' guard matters independently: reduce_loss reports None for it when a microbatch's data processing failed under TWINKLE_FAIL_FAST=0. --- src/twinkle/model/megatron/megatron.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 0b722b57..3ff328b5 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -401,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)) @@ -529,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']