diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 851529c6a..3ff328b53 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 @@ -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'] diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 131ae81cc..04a7bb21d 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/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 88a787ecc..513a3a365 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/processor/test_processor.py b/tests/processor/test_processor.py index d730499d6..a2be36ea5 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.""" diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index 556b8dc83..cff25b4e8 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: