diff --git a/src/paddlefleet/__init__.py b/src/paddlefleet/__init__.py index 46bf4177bb..3d22951a28 100644 --- a/src/paddlefleet/__init__.py +++ b/src/paddlefleet/__init__.py @@ -14,6 +14,10 @@ # Re-export from Paddle for backward compatibility (PaddleFormers imports these from paddlefleet) +from .indexcache_pp_runtime_patch import apply_indexcache_pp_runtime_patch + +apply_indexcache_pp_runtime_patch() + from . import ( parallel_state as parallel_state, training as training, diff --git a/src/paddlefleet/fusions/csa_sparse_attn.py b/src/paddlefleet/fusions/csa_sparse_attn.py index 03d07c22dd..e63ad299eb 100644 --- a/src/paddlefleet/fusions/csa_sparse_attn.py +++ b/src/paddlefleet/fusions/csa_sparse_attn.py @@ -107,6 +107,7 @@ def forward( ctx.query_shape = (b, sq, np_heads, hn) ctx.softmax_scale = float(softmax_scale) ctx.attn_sink_dtype = attn_sink.dtype + ctx.attn_sink_stop_gradient = bool(attn_sink.stop_gradient) ctx.backend = backend query, kv_full, attn_sink, topk_idxs = prepare_inputs( @@ -195,6 +196,7 @@ def backward(ctx, grad_output): ctx.attn_sink_dtype ) + d_attn_sink = None if ctx.attn_sink_stop_gradient else d_attn_sink return (dq, dkv, d_attn_sink, None) diff --git a/src/paddlefleet/indexcache_pp_runtime_patch.py b/src/paddlefleet/indexcache_pp_runtime_patch.py new file mode 100644 index 0000000000..1dc5f23c5d --- /dev/null +++ b/src/paddlefleet/indexcache_pp_runtime_patch.py @@ -0,0 +1,427 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import functools +import os + +import paddle + + +_PATCH_FLAG = "_indexcache_pp_runtime_patch_applied" +_INDEXCACHE_STATE_KEY = "indexcache_state" +_PIPELINE_KEY_ATTR = "_paddlefleet_pipeline_key" +_PIPELINE_SHAPE_ATTR = "_paddlefleet_pipeline_shape" +_PIPELINE_DTYPE_ATTR = "_paddlefleet_pipeline_dtype" + + +def _debug_enabled() -> bool: + return os.environ.get("INDEXCACHE_TRAIN_DEBUG", "0") == "1" + + +def _is_indexcache_key(key) -> bool: + return isinstance(key, str) and key.startswith(_INDEXCACHE_STATE_KEY) + + +def _get_pipeline_key(tensor): + key = getattr(tensor, "key", None) + if key is None: + key = getattr(tensor, _PIPELINE_KEY_ATTR, None) + return key + + +def _save_pipeline_metadata(tensor, key): + setattr(tensor, _PIPELINE_KEY_ATTR, key) + setattr(tensor, _PIPELINE_SHAPE_ATTR, tuple(tensor.shape)) + setattr(tensor, _PIPELINE_DTYPE_ATTR, tensor.dtype) + + +def _has_indexcache_key(tensors) -> bool: + if not isinstance(tensors, (tuple, list)): + tensors = (tensors,) + return any(_is_indexcache_key(_get_pipeline_key(tensor)) for tensor in tensors) + + +def _mark_indexcache_state_stop_gradient(value): + from paddlefleet.transformer.indexcache_state import ( + apply_stop_gradient_mask, + ) + + if value is None: + return value + if not isinstance(value, (tuple, list)): + value = (value,) + return apply_stop_gradient_mask(value) + + +def _describe_tensor_keys(tensors): + if not isinstance(tensors, (tuple, list)): + tensors = (tensors,) + desc = [] + for tensor in tensors: + if not isinstance(tensor, paddle.Tensor): + desc.append(type(tensor).__name__) + continue + desc.append( + { + "key": _get_pipeline_key(tensor), + "shape": list(tensor.shape), + "dtype": str(tensor.dtype), + "stop_gradient": bool(tensor.stop_gradient), + } + ) + return desc + + +def _detach_and_requires_grad(value): + if value is None: + return None + if isinstance(value, dict): + return { + key: _detach_and_requires_grad(item) + for key, item in value.items() + if item is not None + } + if isinstance(value, (tuple, list)): + detached = [_detach_and_requires_grad(item) for item in value] + return tuple(detached) if isinstance(value, tuple) else detached + if isinstance(value, paddle.Tensor): + detached = value.detach() + detached.stop_gradient = value.stop_gradient + return detached + return value + + +def _clone_and_clear_dataptr(value, fake_clone, clear_dataptr=False): + if value is None: + return None + if isinstance(value, dict): + cloned = { + key: _clone_and_clear_dataptr(item, fake_clone, clear_dataptr) + for key, item in value.items() + if item is not None + } + return cloned + if isinstance(value, (tuple, list)): + cloned = [ + _clone_and_clear_dataptr(item, fake_clone, clear_dataptr) + for item in value + if item is not None + ] + return tuple(cloned) if isinstance(value, tuple) else cloned + if isinstance(value, paddle.Tensor): + cloned = fake_clone.apply(value) + cloned.stop_gradient = value.stop_gradient + if clear_dataptr: + cloned._clear_dataptr() + return cloned + return value + + +def _convert_tensor_dict_to_tuple(output_tensor_dict): + output_tensor = [] + for key, tensor in output_tensor_dict.items(): + if tensor is None: + continue + if isinstance(tensor, (list, tuple)): + if key == _INDEXCACHE_STATE_KEY: + tensor = _mark_indexcache_state_stop_gradient(tensor) + for idx, item in enumerate(tensor): + if item is None: + continue + if not isinstance(item, paddle.Tensor): + raise TypeError( + "Pipeline dict tuple values must be paddle.Tensor " + f"or None, but {key}[{idx}] is {type(item).__name__}." + ) + item.key = f"{key} {idx}" + _save_pipeline_metadata(item, item.key) + output_tensor.append(item) + else: + if not isinstance(tensor, paddle.Tensor): + raise TypeError( + "Pipeline dict values must be paddle.Tensor or tensor " + f"tuple/list, but {key} is {type(tensor).__name__}." + ) + tensor.key = key + _save_pipeline_metadata(tensor, key) + output_tensor.append(tensor) + + output_tensor = tuple(output_tensor) + if _debug_enabled() and _has_indexcache_key(output_tensor): + print( + "[INDEXCACHE_PP_FLOW] dict_to_tuple " + f"tensors={_describe_tensor_keys(output_tensor)}", + flush=True, + ) + return output_tensor + + +def _convert_tensor_tuple_to_dict(input_tensor_tuple): + input_tensor_dict = {} + if not isinstance(input_tensor_tuple, tuple): + input_tensor_tuple = (input_tensor_tuple,) + for tensor in input_tensor_tuple: + key = tensor.key + _save_pipeline_metadata(tensor, key) + if " " in key: + real_key, suffix = key.rsplit(" ", 1) + if suffix.isdigit(): + input_tensor_dict.setdefault(real_key, []).append(tensor) + else: + input_tensor_dict[key] = tensor + else: + input_tensor_dict[key] = tensor + delattr(tensor, "key") + if isinstance(input_tensor_dict.get(_INDEXCACHE_STATE_KEY), list): + input_tensor_dict[_INDEXCACHE_STATE_KEY] = tuple( + input_tensor_dict[_INDEXCACHE_STATE_KEY] + ) + if _INDEXCACHE_STATE_KEY in input_tensor_dict: + input_tensor_dict[_INDEXCACHE_STATE_KEY] = ( + _mark_indexcache_state_stop_gradient( + input_tensor_dict[_INDEXCACHE_STATE_KEY] + ) + ) + if _debug_enabled() and _INDEXCACHE_STATE_KEY in input_tensor_dict: + state = input_tensor_dict[_INDEXCACHE_STATE_KEY] + print( + "[INDEXCACHE_PP_FLOW] tuple_to_dict " + f"keys={list(input_tensor_dict.keys())} " + f"state_type={type(state).__name__} " + f"state={_describe_tensor_keys(state)}", + flush=True, + ) + return input_tensor_dict + + +def _tuple_to_dict_helper(input_tensor): + use_dict = False + if isinstance(input_tensor, tuple): + use_dict = len(input_tensor) > 0 and hasattr(input_tensor[0], "key") + else: + use_dict = hasattr(input_tensor, "key") + if use_dict: + input_tensor = _convert_tensor_tuple_to_dict(input_tensor) + return input_tensor, use_dict + + +def _dict_to_tuple_helper(output_tensor): + if isinstance(output_tensor, dict): + return _convert_tensor_dict_to_tuple(output_tensor) + return output_tensor + + +def _collect_input_gradients(inputs): + gradients = [] + zero_filled_keys = [] + for tensor in inputs: + if not isinstance(tensor, paddle.Tensor) or tensor.stop_gradient: + continue + grad = tensor.grad + if grad is None: + key = _get_pipeline_key(tensor) + if not _is_indexcache_key(key): + raise RuntimeError( + "Pipeline input is missing a gradient outside IndexCache " + f"state: key={key!r}, shape={list(tensor.shape)}, " + f"dtype={tensor.dtype}." + ) + grad = paddle.zeros_like(tensor) + grad.stop_gradient = False + zero_filled_keys.append(key) + gradients.append(grad) + if _debug_enabled() and zero_filled_keys: + print( + "[INDEXCACHE_PP_GRAD] zero_filled_keys=" + f"{zero_filled_keys}", + flush=True, + ) + return tuple(gradients) + + +def _zeros_from_tensor_metadata(tensor): + shape = getattr(tensor, _PIPELINE_SHAPE_ATTR, None) + dtype = getattr(tensor, _PIPELINE_DTYPE_ATTR, None) + if shape is None or dtype is None: + raise RuntimeError( + "IndexCache pipeline tensor lacks preserved shape/dtype metadata: " + f"key={_get_pipeline_key(tensor)!r}." + ) + grad = paddle.zeros( + shape=list(shape), + dtype=dtype, + ) + grad.stop_gradient = False + return grad + + +def _normalize_pipeline_input_gradients(input_tensor, input_tensor_grad): + if input_tensor is None: + return input_tensor_grad + + is_tuple_input = isinstance(input_tensor, tuple) + inputs = input_tensor if is_tuple_input else (input_tensor,) + if not _has_indexcache_key(inputs): + return input_tensor_grad + + differentiable_inputs = [ + tensor + for tensor in inputs + if isinstance(tensor, paddle.Tensor) and not tensor.stop_gradient + ] + if is_tuple_input: + if not isinstance(input_tensor_grad, (tuple, list)): + raise RuntimeError( + "IndexCache pipeline input gradients must preserve tuple " + f"structure, but got {type(input_tensor_grad).__name__}." + ) + gradients = list(input_tensor_grad) + else: + gradients = [input_tensor_grad] + + if len(gradients) != len(differentiable_inputs): + raise RuntimeError( + "IndexCache pipeline input gradient arity mismatch: " + f"inputs={len(differentiable_inputs)}, gradients={len(gradients)}." + ) + + zero_filled_keys = [] + for idx, (tensor, grad) in enumerate(zip(differentiable_inputs, gradients)): + key = _get_pipeline_key(tensor) + if grad is None: + if not _is_indexcache_key(key): + raise RuntimeError( + "Pipeline input is missing a gradient outside IndexCache " + f"state: key={key!r}, shape={list(tensor.shape)}, " + f"dtype={tensor.dtype}." + ) + grad = _zeros_from_tensor_metadata(tensor) + gradients[idx] = grad + zero_filled_keys.append(key) + elif not isinstance(grad, paddle.Tensor): + raise TypeError( + "Pipeline input gradients must be paddle.Tensor or None, " + f"but key={key!r} has {type(grad).__name__}." + ) + + if _debug_enabled() and zero_filled_keys: + print( + "[INDEXCACHE_PP_GRAD] boundary=pipeline zero_filled_keys=" + f"{zero_filled_keys}", + flush=True, + ) + return tuple(gradients) if is_tuple_input else gradients[0] + + +def _wrap_pipeline_backward_step(original_backward_step): + @functools.wraps(original_backward_step) + def backward_step(self, input_tensor, *args, **kwargs): + input_tensor_grad = original_backward_step( + self, + input_tensor, + *args, + **kwargs, + ) + return _normalize_pipeline_input_gradients( + input_tensor, + input_tensor_grad, + ) + + return backward_step + + +def _schedule_node_backward(self, output_grad=None, scaler=None): + if output_grad is None: + if isinstance(self.outputs, (tuple, list)): + assert len(self.outputs) == 1 + outputs = self.outputs[0] + else: + outputs = self.outputs + assert isinstance(outputs, paddle.Tensor) + if scaler is not None: + paddle.autograd.backward(scaler.scale(outputs)) + else: + paddle.autograd.backward(outputs) + else: + is_output_grad_tuple = isinstance(output_grad, tuple) + if not isinstance(output_grad, (tuple, list)): + is_output_grad_tuple = True + output_grad = (output_grad,) + + outputs = _dict_to_tuple_helper(self.outputs) + if not isinstance(outputs, (tuple, list)): + outputs = (outputs,) + outputs = [ + tensor + for tensor in outputs + if isinstance(tensor, paddle.Tensor) and not tensor.stop_gradient + ] + + output_grad = [grad for grad in output_grad if grad is not None] + output_grad = ( + tuple(output_grad) if is_output_grad_tuple else list(output_grad) + ) + + assert len(outputs) == len(output_grad), ( + f"{len(outputs)} of {type(outputs[0])} vs " + f"{len(output_grad)} of {type(output_grad[0])}" + ) + paddle.autograd.backward(outputs, output_grad) + + inputs = _dict_to_tuple_helper(self.inputs) + if not isinstance(inputs, (tuple, list)): + inputs = (inputs,) + grad = _collect_input_gradients(inputs) + self._reset_states() + return grad + + +def apply_indexcache_pp_runtime_patch() -> None: + import paddle.distributed.fleet.meta_parallel as meta_parallel + import paddle.distributed.fleet.meta_parallel.pipeline_parallel as pipeline_parallel + import paddle.distributed.fleet.meta_parallel.pp_utils.forward_backward_overlap_utils as fbo + import paddle.distributed.fleet.meta_parallel.pp_utils.utils as pp_utils + + if getattr(pp_utils, _PATCH_FLAG, False): + return + + def clone_and_clear_dataptr(outputs, clear_dataptr=False): + return _clone_and_clear_dataptr( + outputs, + fbo.FakeClone, + clear_dataptr=clear_dataptr, + ) + + pp_utils.convert_tensor_dict_to_tuple = _convert_tensor_dict_to_tuple + pp_utils.convert_tensor_tuple_to_dict = _convert_tensor_tuple_to_dict + pp_utils.tuple_to_dict_helper = _tuple_to_dict_helper + pp_utils.dict_to_tuple_helper = _dict_to_tuple_helper + setattr(pp_utils, _PATCH_FLAG, True) + + meta_parallel.dict_to_tuple_helper = _dict_to_tuple_helper + meta_parallel.tuple_to_dict_helper = _tuple_to_dict_helper + + fbo.detach_and_requires_grad = _detach_and_requires_grad + fbo.clone_and_clear_dataptr = clone_and_clear_dataptr + fbo.ScheduleNode.backward = _schedule_node_backward + setattr(fbo, _PATCH_FLAG, True) + + pipeline_parallel.PipelineParallel._backward_step = ( + _wrap_pipeline_backward_step( + pipeline_parallel.PipelineParallel._backward_step + ) + ) + setattr(pipeline_parallel, _PATCH_FLAG, True) diff --git a/src/paddlefleet/models/gpt/gpt_embedding.py b/src/paddlefleet/models/gpt/gpt_embedding.py index 345ceb81fa..ff1f4df77b 100644 --- a/src/paddlefleet/models/gpt/gpt_embedding.py +++ b/src/paddlefleet/models/gpt/gpt_embedding.py @@ -216,6 +216,11 @@ def forward( decoder_input, text_padding_indices, 0 ) input_ids_for_moe_mask = input_ids + if ( + input_ids_for_moe_mask is None + and getattr(self.config, "moe_n_hash_layers", 0) > 0 + ): + input_ids_for_moe_mask = input_ids if ( self.config.num_nextn_predict_layers is not None and self.config.num_nextn_predict_layers > 0 diff --git a/src/paddlefleet/transformer/csa_attention.py b/src/paddlefleet/transformer/csa_attention.py index a0a684c173..86c82369c0 100644 --- a/src/paddlefleet/transformer/csa_attention.py +++ b/src/paddlefleet/transformer/csa_attention.py @@ -51,6 +51,24 @@ fused_qk_topk_naive, rotate_activation, ) +from paddlefleet.transformer.indexcache_state import ( + INDEXCACHE_DISTILL_STATE_K, + INDEXCACHE_DISTILL_STATE_PRODUCER_LAYER, + INDEXCACHE_DISTILL_STATE_Q, + INDEXCACHE_DISTILL_STATE_SERVED_COUNT, + INDEXCACHE_DISTILL_STATE_TOPK_INDICES, + INDEXCACHE_DISTILL_STATE_TOPK_PROBS, + INDEXCACHE_DISTILL_STATE_WEIGHTS, + INDEXCACHE_STATE_KIND_DISTILL, + INDEXCACHE_STATE_KIND_INVALID, + INDEXCACHE_STATE_KIND_TOPK_ONLY, + INDEXCACHE_STATE_TOPK_IDXS, + INDEXCACHE_TOPK_ONLY_STATE_LEN, + INDEXCACHE_TOPK_ONLY_STATE_PRODUCER_LAYER, + apply_stop_gradient_mask, + detach_stop_gradient_tensor, + state_kind, +) from paddlefleet.transformer.utils import ( get_doc_lens, get_doc_starts, @@ -72,7 +90,6 @@ map_compressed_topk_to_kv_full_cp, ) - class LinearBF16FP32Func(paddle.autograd.PyLayer): """BF16 activation x BF16 weight -> FP32 output autograd function. @@ -800,6 +817,24 @@ def _compute_fused_csa_indexer_loss_forward( return loss, topk_indices, topk_probs, target +def _compute_csa_selected_set_kl_loss( + target: Tensor, + topk_probs: Tensor, + loss_coeff: float, + loss_mask: Tensor | None = None, + global_valid_count: float | None = None, +) -> Tensor: + eps = 1e-10 + kl_per_elem = target * ( + paddle.log(target + eps) - paddle.log(topk_probs + eps) + ) + kl_per_pos = kl_per_elem.sum(axis=-1) + if loss_mask is not None: + lm = loss_mask.reshape(kl_per_pos.shape).astype(kl_per_pos.dtype) + return (kl_per_pos * lm).sum() / global_valid_count * float(loss_coeff) + return kl_per_pos.mean() * float(loss_coeff) + + class TileLangCSAIndexerLossAutoScaler(paddle.autograd.PyLayer): """Attach TileLang CSA indexer loss gradients to the main output. @@ -823,6 +858,12 @@ def forward( num_rows_override: float | None = None, loss_mask: Tensor | None = None, ) -> Tensor: + ctx.input_stop_gradients = ( + bool(output.stop_gradient), + bool(index_q.stop_gradient), + bool(weights.stop_gradient), + bool(index_k_comp.stop_gradient), + ) ctx.save_for_backward( index_q.detach(), weights.detach(), @@ -934,7 +975,35 @@ def backward(ctx, grad_output: Tensor): if grad_k.dtype != index_k_comp.dtype: grad_k = grad_k.cast(index_k_comp.dtype) - grads = (grad_output, grad_q, grad_weights, grad_k) + (None,) * ( + if os.environ.get("INDEXCACHE_TRAIN_DEBUG", "0") == "1": + q_norm = float( + paddle.linalg.norm(grad_q.cast(paddle.float32)).item() + ) + weights_norm = float( + paddle.linalg.norm(grad_weights.cast(paddle.float32)).item() + ) + k_norm = float( + paddle.linalg.norm(grad_k.cast(paddle.float32)).item() + ) + print( + "[INDEXCACHE_DISTILL_GRAD] " + f"backend={ctx.indexer_backend} " + f"loss_coeff={ctx.loss_coeff:.8g} " + f"input_stop_gradients={ctx.input_stop_gradients} " + f"q_grad_norm={q_norm:.8g} " + f"weights_grad_norm={weights_norm:.8g} " + f"k_grad_norm={k_norm:.8g}", + flush=True, + ) + + grad_slots = [grad_output, grad_q, grad_weights, grad_k] + for idx, stop_gradient in enumerate( + getattr(ctx, "input_stop_gradients", ()) + ): + if stop_gradient: + grad_slots[idx] = None + + grads = tuple(grad_slots) + (None,) * ( 4 if getattr(ctx, "loss_mask", None) is not None else 3 ) return grads @@ -1578,6 +1647,510 @@ def __init__( else: self.indexer = None + def _indexcache_pattern(self) -> str | None: + pattern = getattr(self.config, "index_topk_pattern", None) + if pattern is None: + return None + pattern = str(pattern).strip().upper() + if not pattern: + return None + invalid_chars = sorted(set(pattern) - {"F", "S"}) + if invalid_chars: + raise ValueError( + "index_topk_pattern may only contain 'F' and 'S', " + f"got invalid chars: {invalid_chars}." + ) + if pattern[0] != "F": + raise ValueError("index_topk_pattern must start with 'F'.") + return pattern + + def _indexcache_recompute_enabled(self) -> bool: + return bool(getattr(self.config, "recompute_granularity", None)) + + def _indexcache_in_recompute(self) -> bool: + return ( + self._indexcache_recompute_enabled() + and self.training + and not paddle.is_grad_enabled() + ) + + def _indexcache_context_msg( + self, c4_ordinal: int, pattern: str + ) -> str: + return ( + f"layer_number={self.layer_number}, " + f"c4_ordinal={c4_ordinal}, " + f"pattern={pattern}, " + "recompute_granularity=" + f"{getattr(self.config, 'recompute_granularity', None)}, " + "pipeline_model_parallel_size=" + f"{getattr(self.config, 'pipeline_model_parallel_size', 1)}, " + "context_parallel_size=" + f"{getattr(self.config, 'context_parallel_size', 1)}" + ) + + def _indexcache_requires_explicit_state(self) -> bool: + pp_size = int( + getattr(self.config, "pipeline_model_parallel_size", 1) or 1 + ) + cp_size = int(getattr(self.config, "context_parallel_size", 1) or 1) + return self._indexcache_recompute_enabled() or pp_size > 1 or cp_size > 1 + + @staticmethod + def _indexcache_state_kind(indexcache_state: tuple | list | None) -> str: + return state_kind(indexcache_state) + + def _indexcache_validate_state_for_reuse( + self, + indexcache_state: tuple | list | None, + c4_ordinal: int, + pattern: str, + ) -> str: + state_kind = self._indexcache_state_kind(indexcache_state) + if state_kind == INDEXCACHE_STATE_KIND_INVALID: + raise ValueError( + "IndexCache state must be either topk-only " + f"({INDEXCACHE_TOPK_ONLY_STATE_LEN} tensors) or distill " + "(8 tensors), got " + f"len={len(indexcache_state)}. " + + self._indexcache_context_msg(c4_ordinal, pattern) + ) + if ( + state_kind == INDEXCACHE_STATE_KIND_DISTILL + and not self._indexcache_multi_layer_distill_enabled() + ): + raise RuntimeError( + "IndexCache reuse-only expects a topk-only indexcache_state; " + "distill-state tensors require " + "indexcache_multi_layer_distill=True. " + f"state_kind={state_kind}. " + + self._indexcache_context_msg(c4_ordinal, pattern) + ) + return state_kind + + def _indexcache_debug(self, msg: str) -> None: + if os.environ.get("INDEXCACHE_TRAIN_DEBUG", "0") == "1": + cp_msg = ( + f" cp_rank={self.cp_rank} cp_size={self.cp_size}" + if self.cp_enabled + else "" + ) + recompute_msg = ( + " recompute_enabled=" + f"{self._indexcache_recompute_enabled()}" + " in_recompute=" + f"{self._indexcache_in_recompute()}" + " grad_enabled=" + f"{paddle.is_grad_enabled()}" + ) + print( + f"[INDEXCACHE_TRAIN] layer={self.layer_number}{cp_msg}" + f"{recompute_msg} {msg}", + flush=True, + ) + + def _indexcache_distill_debug(self, msg: str) -> None: + if os.environ.get("INDEXCACHE_TRAIN_DEBUG", "0") == "1": + cp_msg = ( + f" cp_rank={self.cp_rank} cp_size={self.cp_size}" + if self.cp_enabled + else "" + ) + print( + f"[INDEXCACHE_DISTILL] layer={self.layer_number}{cp_msg} {msg}", + flush=True, + ) + + def _indexcache_multi_layer_distill_enabled(self) -> bool: + return bool(getattr(self.config, "indexcache_multi_layer_distill", False)) + + @staticmethod + def _indexcache_served_count(pattern: str, c4_ordinal: int) -> int: + next_f = pattern.find("F", c4_ordinal + 1) + if next_f == -1: + next_f = len(pattern) + return max(1, next_f - c4_ordinal) + + def _indexcache_scaled_loss_coeff(self, served_count: int) -> float: + coeff = getattr(self.config, "dsa_indexer_loss_coeff", 0.0) or 0.0 + return float(coeff) / float(max(int(served_count), 1)) + + def _indexcache_clear_cached_state(self) -> None: + self.config._indexcache_last_topk_idxs = None + self.config._indexcache_last_layer_number = None + self.config._indexcache_last_distill_state = None + self.config._indexcache_last_served_count = None + + def _indexcache_c4_layers(self) -> list[int]: + ratios = getattr(self.config, "csa_compress_ratios", None) or [] + empty_head = int( + getattr(self.config, "num_empty_layers_add_in_head", 0) or 0 + ) + return [ + int(layer_idx) + empty_head + for layer_idx, ratio in enumerate(ratios) + if int(ratio) == 4 + ] + + def _indexcache_infer_producer_layer( + self, c4_ordinal: int, pattern: str + ) -> int | None: + producer_ordinal = min(c4_ordinal - 1, len(pattern) - 1) + while producer_ordinal >= 0 and pattern[producer_ordinal] != "F": + producer_ordinal -= 1 + if producer_ordinal < 0: + return None + + c4_layers = self._indexcache_c4_layers() + if producer_ordinal >= len(c4_layers): + return None + return c4_layers[producer_ordinal] + + @staticmethod + def _indexcache_has_future_served_layer( + pattern: str, c4_ordinal: int + ) -> bool: + next_f = pattern.find("F", c4_ordinal + 1) + end = next_f if next_f != -1 else len(pattern) + return "S" in pattern[c4_ordinal + 1 : end] + + def _indexcache_next_c4_action(self) -> tuple[int, str, str] | None: + pattern = self._indexcache_pattern() + if pattern is None or self.compress_ratio != 4 or self.indexer is None: + return None + + c4_layers = self._indexcache_c4_layers() + if self.layer_number not in c4_layers: + raise RuntimeError( + "IndexCache C4 static mapping could not find the current " + f"layer={self.layer_number} in csa_compress_ratios C4 layers " + f"{c4_layers}." + ) + c4_ordinal = c4_layers.index(self.layer_number) + if c4_ordinal >= len(pattern): + raise ValueError( + "index_topk_pattern must cover every C4 layer in this " + f"configuration. pattern={pattern}, c4_layers={c4_layers}, " + f"current_c4_ordinal={c4_ordinal}." + ) + + if c4_ordinal == 0: + self._indexcache_clear_cached_state() + + action = pattern[c4_ordinal] + return c4_ordinal, action, pattern + + @staticmethod + def _indexcache_tensor_to_int(value: Tensor | int | None) -> int | None: + if value is None: + return None + if isinstance(value, int): + return value + try: + return int(value.item()) + except Exception: + return int(value.numpy().reshape([-1])[0]) + + @staticmethod + def _indexcache_forward_state_tensor(value: Tensor) -> Tensor: + return detach_stop_gradient_tensor(value) + + def _indexcache_pack_state( + self, + compress_topk_idxs: Tensor, + tilelang_indexer_loss_state: tuple | None = None, + served_count: int | None = None, + ) -> tuple[Tensor, ...]: + state = [self._indexcache_forward_state_tensor(compress_topk_idxs)] + include_distill_state = ( + self._indexcache_multi_layer_distill_enabled() + and tilelang_indexer_loss_state is not None + ) + if include_distill_state: + ( + q_indexer_bf, + weights_indexer_bf, + k_indexer_bf, + topk_indices_compressed, + topk_probs, + *_, + ) = tilelang_indexer_loss_state + state.extend( + [ + q_indexer_bf, + weights_indexer_bf, + k_indexer_bf, + self._indexcache_forward_state_tensor( + topk_indices_compressed + ), + self._indexcache_forward_state_tensor(topk_probs), + self._indexcache_forward_state_tensor( + paddle.full([1], self.layer_number, dtype="int64") + ), + self._indexcache_forward_state_tensor( + paddle.full([1], int(served_count or 1), dtype="int64") + ), + ] + ) + else: + state.extend( + [ + self._indexcache_forward_state_tensor( + paddle.full([1], self.layer_number, dtype="int64") + ), + self._indexcache_forward_state_tensor( + paddle.full([1], int(served_count or 1), dtype="int64") + ), + ] + ) + return apply_stop_gradient_mask(tuple(state)) + + def _indexcache_state_topk( + self, + indexcache_state: tuple | list | None, + c4_ordinal: int, + pattern: str, + ) -> tuple[Tensor | None, int | None, str]: + state_kind = self._indexcache_validate_state_for_reuse( + indexcache_state, c4_ordinal, pattern + ) + if not indexcache_state: + return None, None, state_kind + producer_layer = None + if state_kind == INDEXCACHE_STATE_KIND_DISTILL: + producer_layer = self._indexcache_tensor_to_int( + indexcache_state[INDEXCACHE_DISTILL_STATE_PRODUCER_LAYER] + ) + elif state_kind == INDEXCACHE_STATE_KIND_TOPK_ONLY: + producer_layer = self._indexcache_tensor_to_int( + indexcache_state[INDEXCACHE_TOPK_ONLY_STATE_PRODUCER_LAYER] + ) + if producer_layer is None: + producer_layer = getattr( + self.config, "_indexcache_last_layer_number", None + ) + return ( + indexcache_state[INDEXCACHE_STATE_TOPK_IDXS], + producer_layer, + state_kind, + ) + + def _indexcache_cache_topk( + self, + compress_topk_idxs: Tensor, + c4_ordinal: int, + pattern: str, + tilelang_indexer_loss_state: tuple | None = None, + served_count: int | None = None, + loss_scale: float | None = None, + ) -> tuple[Tensor, ...]: + self.config._indexcache_last_topk_idxs = compress_topk_idxs.detach() + self.config._indexcache_last_layer_number = self.layer_number + if ( + self._indexcache_multi_layer_distill_enabled() + and tilelang_indexer_loss_state is not None + ): + state_kind = INDEXCACHE_STATE_KIND_DISTILL + self.config._indexcache_last_distill_state = tilelang_indexer_loss_state + self.config._indexcache_last_served_count = served_count + self._indexcache_distill_debug( + "action=producer " + f"c4_ordinal={c4_ordinal} served_count={served_count} " + f"loss_scale={loss_scale:.8g}" + ) + else: + state_kind = INDEXCACHE_STATE_KIND_TOPK_ONLY + self._indexcache_debug( + "action=produce " + f"c4_ordinal={c4_ordinal} pattern={pattern} " + f"state_kind={state_kind} " + f"topk_shape={list(compress_topk_idxs.shape)}" + ) + return self._indexcache_pack_state( + compress_topk_idxs, + tilelang_indexer_loss_state, + served_count, + ) + + def _indexcache_reuse_topk( + self, + b: int, + sq: int, + c4_ordinal: int, + pattern: str, + indexcache_state: tuple | list | None = None, + ) -> Tensor: + cached, producer_layer, state_kind = self._indexcache_state_topk( + indexcache_state, c4_ordinal, pattern + ) + if cached is None and not self._indexcache_requires_explicit_state(): + cached = getattr(self.config, "_indexcache_last_topk_idxs", None) + producer_layer = getattr( + self.config, "_indexcache_last_layer_number", None + ) + if cached is not None: + state_kind = "config_fallback" + if cached is None: + raise RuntimeError( + "index_topk_pattern requested reuse before an explicit " + "producer top-k state exists. " + f"state_kind={state_kind}. " + + self._indexcache_context_msg(c4_ordinal, pattern) + ) + cached_shape = list(cached.shape) + if len(cached_shape) < 3 or cached_shape[0] != b or cached_shape[1] != sq: + raise ValueError( + "Cached IndexCache top-k shape does not match the current " + f"layer input. cached={cached_shape}, current_batch={b}, " + f"current_seq={sq}." + ) + if producer_layer is None: + producer_layer = self._indexcache_infer_producer_layer( + c4_ordinal, pattern + ) + self._indexcache_debug( + "action=reuse " + f"c4_ordinal={c4_ordinal} pattern={pattern} " + f"producer_layer={producer_layer} state_kind={state_kind} " + f"topk_shape={cached_shape}" + ) + return cached + + def _indexcache_served_distill_state( + self, + query: Tensor, + compressed_kv: Tensor, + c4_ordinal: int, + pattern: str, + loss_mask: Tensor | None = None, + global_valid_count: float | None = None, + indexcache_state: tuple | list | None = None, + ) -> tuple | None: + if not ( + self._indexcache_multi_layer_distill_enabled() + and self.training + and paddle.is_grad_enabled() + ): + return None + + producer_state = None + producer_layer = None + served_count = None + if indexcache_state is not None: + state_kind = self._indexcache_state_kind(indexcache_state) + if state_kind == INDEXCACHE_STATE_KIND_DISTILL: + producer_state = ( + indexcache_state[INDEXCACHE_DISTILL_STATE_Q], + indexcache_state[INDEXCACHE_DISTILL_STATE_WEIGHTS], + indexcache_state[INDEXCACHE_DISTILL_STATE_K], + indexcache_state[INDEXCACHE_DISTILL_STATE_TOPK_INDICES], + indexcache_state[INDEXCACHE_DISTILL_STATE_TOPK_PROBS], + None, + None, + getattr(self.config, "csa_indexer_backend", "tilelang"), + None, + None, + ) + producer_layer = self._indexcache_tensor_to_int( + indexcache_state[INDEXCACHE_DISTILL_STATE_PRODUCER_LAYER] + ) + served_count = self._indexcache_tensor_to_int( + indexcache_state[INDEXCACHE_DISTILL_STATE_SERVED_COUNT] + ) + else: + raise RuntimeError( + "indexcache_multi_layer_distill received a producer " + "top-k state without producer loss tensors. " + f"c4_ordinal={c4_ordinal}, pattern={pattern}." + ) + + if producer_state is None and not self._indexcache_requires_explicit_state(): + producer_state = getattr( + self.config, "_indexcache_last_distill_state", None + ) + producer_layer = getattr( + self.config, "_indexcache_last_layer_number", None + ) + served_count = getattr( + self.config, "_indexcache_last_served_count", None + ) + if producer_state is None or served_count is None: + raise RuntimeError( + "indexcache_multi_layer_distill requested an S-layer target " + "before an explicit producer distill state exists. " + + self._indexcache_context_msg(c4_ordinal, pattern) + ) + + ( + q_indexer_bf, + weights_indexer_bf, + k_indexer_bf, + topk_indices_compressed, + topk_probs, + _producer_target, + _producer_loss_coeff, + backend, + _producer_num_rows, + _producer_loss_mask, + ) = producer_state + key_comp_mla = compressed_kv.detach() + + if self.tp_group is not None and getattr(self.tp_group, "nranks", 1) > 1: + target = _compute_attn_target_on_selected_set( + query.detach(), + key_comp_mla, + topk_indices_compressed.detach(), + self.softmax_scale, + self.tp_group, + ) + else: + from paddlefleet.tilelang_ops import csa_attn_target_reducesum + + target = csa_attn_target_reducesum( + query.detach(), + key_comp_mla, + topk_indices_compressed.detach(), + self.softmax_scale, + ) + + loss_scale = self._indexcache_scaled_loss_coeff(int(served_count)) + if self.cp_enabled and loss_mask is None: + loss_scale /= float(self.cp_size) + if loss_scale > 0: + loss = _compute_csa_selected_set_kl_loss( + target.detach(), + topk_probs.detach(), + loss_scale, + loss_mask, + global_valid_count, + ) + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=loss, + layer_number=self.layer_number, + num_layers=DSAIndexerLossLoggingHelper.get_total_num_layers( + self.config + ), + ) + self._indexcache_distill_debug( + "action=served " + f"c4_ordinal={c4_ordinal} pattern={pattern} " + f"producer_layer={producer_layer} served_count={served_count} " + f"topk_shape={list(topk_indices_compressed.shape)} " + f"loss_scale={loss_scale:.8g}" + ) + return ( + q_indexer_bf, + weights_indexer_bf, + k_indexer_bf, + topk_indices_compressed, + topk_probs, + target, + loss_scale, + backend, + global_valid_count if loss_mask is not None else None, + loss_mask, + ) + def _compute_indexer_compressed_topk_idxs( self, query: Tensor, @@ -1589,6 +2162,8 @@ def _compute_indexer_compressed_topk_idxs( startend_row_indices: Tensor | None = None, loss_mask: Tensor | None = None, global_valid_count: float | None = None, + indexer_loss_coeff_override: float | None = None, + materialize_distill_state: bool = False, ) -> tuple[Tensor, Tensor | None, tuple | None]: """Build indexer-selected compressed KV indices and loss state.""" b, sq, np_heads, _ = query.shape @@ -1608,10 +2183,16 @@ def _compute_indexer_compressed_topk_idxs( indexer_backend = getattr( self.config, "csa_indexer_backend", "tilelang" ) + if materialize_distill_state and indexer_backend != "tilelang": + raise NotImplementedError( + "IndexCache distill recompute materialization currently " + "supports only csa_indexer_backend='tilelang'." + ) # The indexer loss path is only active during the grad-enabled forward. # Full recompute runs the first forward under no_grad; that pass should - # only materialize main-attention indices. The backend branch remains - # fixed across both forwards. + # materialize main-attention indices, plus producer distill state when + # a later S layer will consume it. The backend branch remains fixed + # across both forwards. need_indexer_loss = self.training and paddle.is_grad_enabled() loss_topk_effective = _resolve_csa_indexer_loss_topk_effective( self.config, @@ -1642,8 +2223,10 @@ def compute_fused_indexer_loss(backend: str): # Training grad-enabled forward with TileLang/cuDNN backend. # The same backend's top-k-only kernel is used during recompute's # first no-grad forward below. - indexer_loss_coeff = getattr( - self.config, "dsa_indexer_loss_coeff", 0.0 + indexer_loss_coeff = ( + indexer_loss_coeff_override + if indexer_loss_coeff_override is not None + else getattr(self.config, "dsa_indexer_loss_coeff", 0.0) ) q_indexer_bf, k_indexer_bf, weights_indexer_bf = ( self.indexer.forward_before_topk( @@ -1732,7 +2315,7 @@ def compute_fused_indexer_loss(backend: str): topk_indices_compressed, tilelang_indexer_loss_state, ) = compute_fused_indexer_loss("tilelang") - else: # First recompute no-grad forward; only materialize TileLang top-k for attention. + else: # First recompute no-grad forward; materialize attention top-k and optional producer distill state. from paddlefleet.tilelang_ops import csa_indexer_topk_fwd with paddle.no_grad(): @@ -1741,15 +2324,40 @@ def compute_fused_indexer_loss(backend: str): x_det, qr_det, startend_row_indices ) ) - tl_topk_indices, _tl_topk_scores = csa_indexer_topk_fwd( + tl_topk_effective = ( + loss_topk_effective + if materialize_distill_state + else attn_topk_effective + ) + tl_topk_indices, tl_topk_scores = csa_indexer_topk_fwd( q_indexer_tl, k_indexer_tl, weights_indexer_tl, ratio=self.compress_ratio, - topk_effective=attn_topk_effective, + topk_effective=tl_topk_effective, valid_range=valid_range, ) topk_indices_compressed = tl_topk_indices + if materialize_distill_state: + indexer_loss_coeff = ( + indexer_loss_coeff_override + if indexer_loss_coeff_override is not None + else getattr( + self.config, "dsa_indexer_loss_coeff", 0.0 + ) + ) + tilelang_indexer_loss_state = ( + q_indexer_tl, + weights_indexer_tl, + k_indexer_tl, + topk_indices_compressed, + tl_topk_scores, + None, + float(indexer_loss_coeff), + "tilelang", + global_valid_count if loss_mask is not None else None, + loss_mask, + ) elif ( indexer_backend == "unfused" @@ -1760,8 +2368,10 @@ def compute_fused_indexer_loss(backend: str): x_det, qr_det, startend_row_indices ) ) - indexer_loss_coeff = getattr( - self.config, "dsa_indexer_loss_coeff", 0.0 + indexer_loss_coeff = ( + indexer_loss_coeff_override + if indexer_loss_coeff_override is not None + else getattr(self.config, "dsa_indexer_loss_coeff", 0.0) ) key_for_loss = compressed_kv.unsqueeze(2).expand( [-1, -1, np_heads, -1] @@ -1829,6 +2439,7 @@ def forward( x: Tensor = None, qr: Tensor = None, input_ids: Tensor | None = None, + indexcache_state: tuple | list | None = None, ) -> Tensor: """Forward pass for CompressedSparseAttention. @@ -1844,6 +2455,8 @@ def forward( output: [b, sq, np * v_head_dim] """ b, sq, np_heads, hn = query.shape + indexcache_state_next = indexcache_state + indexcache_state_updated = False if startend_row_indices is not None: assert b == 1, ( @@ -1889,7 +2502,7 @@ def forward( global_valid_count = None if self.cp_enabled: - return self._forward_cp( + cp_result = self._forward_cp( query, key, x, @@ -1897,7 +2510,18 @@ def forward( startend_row_indices, loss_mask=loss_mask, global_valid_count=global_valid_count, + indexcache_state=indexcache_state, ) + cp_indexcache_state_updated = isinstance(cp_result, tuple) + if isinstance(cp_result, tuple): + output, indexcache_state_next = cp_result + else: + output = cp_result + if indexcache_state_next is not None: + return output, indexcache_state_next + if cp_indexcache_state_updated: + return output, None + return output if startend_row_indices is not None and self.compress_ratio > 1: doc_lens = get_doc_lens(startend_row_indices) @@ -1941,6 +2565,7 @@ def forward( # Step 4: Compressed indices indexer_loss = None tilelang_indexer_loss_state = None + indexcache_served_loss_state = None if ( self.compress_ratio > 1 @@ -1948,21 +2573,93 @@ def forward( and actual_n_compressed > 0 ): if self.indexer is not None: - ( - compress_topk_idxs, - indexer_loss, - tilelang_indexer_loss_state, - ) = self._compute_indexer_compressed_topk_idxs( - query, - x, - qr, - compressed_kv, - n_compressed, - offset, - startend_row_indices, - loss_mask=loss_mask, - global_valid_count=global_valid_count, - ) + indexcache_action = self._indexcache_next_c4_action() + if ( + indexcache_action is not None + and indexcache_action[1] == "S" + ): + c4_ordinal, _action, pattern = indexcache_action + compress_topk_idxs = self._indexcache_reuse_topk( + b, + sq, + c4_ordinal, + pattern, + indexcache_state=indexcache_state, + ) + indexcache_served_loss_state = ( + self._indexcache_served_distill_state( + query, + compressed_kv, + c4_ordinal, + pattern, + loss_mask=loss_mask, + global_valid_count=global_valid_count, + indexcache_state=indexcache_state, + ) + ) + if self._indexcache_has_future_served_layer( + pattern, c4_ordinal + ): + indexcache_state_next = indexcache_state + else: + indexcache_state_next = None + self._indexcache_clear_cached_state() + indexcache_state_updated = True + else: + served_count = None + loss_coeff_override = None + materialize_distill_state = False + if ( + indexcache_action is not None + and self._indexcache_multi_layer_distill_enabled() + ): + c4_ordinal, _action, pattern = indexcache_action + served_count = self._indexcache_served_count( + pattern, c4_ordinal + ) + loss_coeff_override = self._indexcache_scaled_loss_coeff( + served_count + ) + materialize_distill_state = ( + self._indexcache_has_future_served_layer( + pattern, c4_ordinal + ) + ) + ( + compress_topk_idxs, + indexer_loss, + tilelang_indexer_loss_state, + ) = self._compute_indexer_compressed_topk_idxs( + query, + x, + qr, + compressed_kv, + n_compressed, + offset, + startend_row_indices, + loss_mask=loss_mask, + global_valid_count=global_valid_count, + indexer_loss_coeff_override=loss_coeff_override, + materialize_distill_state=materialize_distill_state, + ) + if indexcache_action is not None: + c4_ordinal, _action, pattern = indexcache_action + produced_state = self._indexcache_cache_topk( + compress_topk_idxs, + c4_ordinal, + pattern, + tilelang_indexer_loss_state, + served_count, + loss_coeff_override, + ) + if self._indexcache_has_future_served_layer( + pattern, c4_ordinal + ): + indexcache_state_next = produced_state + else: + indexcache_state_next = None + self._indexcache_clear_cached_state() + indexcache_state_updated = True else: # ratio=128: attend to all compressed positions compress_topk_idxs = get_compress_topk_idxs( @@ -1993,14 +2690,31 @@ def forward( ) # Step 6: Attach indexer loss - if tilelang_indexer_loss_state is not None and self.training: + if ( + tilelang_indexer_loss_state is not None + and self.training + and paddle.is_grad_enabled() + ): output = TileLangCSAIndexerLossAutoScaler.apply( output, *tilelang_indexer_loss_state, ) - elif indexer_loss is not None and self.training: + elif ( + indexer_loss is not None + and self.training + and paddle.is_grad_enabled() + ): output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + if indexcache_served_loss_state is not None: + output = TileLangCSAIndexerLossAutoScaler.apply( + output, + *indexcache_served_loss_state, + ) + if indexcache_state_next is not None: + return output, indexcache_state_next + if indexcache_state_updated: + return output, None return output def _forward_cp( @@ -2012,6 +2726,7 @@ def _forward_cp( startend_row_indices: Tensor | None = None, loss_mask: Tensor | None = None, global_valid_count: float | None = None, + indexcache_state: tuple | list | None = None, ) -> Tensor: """CP-aware forward: local compress + all-gather, sparse attention. @@ -2028,6 +2743,8 @@ def _forward_cp( reduce_scatter(SUM) + optimizer x cp_size aggregates them correctly """ b, sq, np_heads, hn = query.shape + indexcache_state_next = indexcache_state + indexcache_state_updated = False sq_global = sq * self.cp_size position_offset = self.cp_rank * sq q_positions = paddle.arange( @@ -2096,6 +2813,7 @@ def _forward_cp( # Step 3: Compressed topk + optional fused indexer loss indexer_loss = None tilelang_indexer_loss_state = None + indexcache_served_loss_state = None if ( self.compress_ratio > 1 @@ -2103,233 +2821,349 @@ def _forward_cp( and actual_n_compressed > 0 ): if self.indexer is not None: - x_det = x.detach() - qr_det = qr.detach() - if self.training: - x_det.stop_gradient = False - qr_det.stop_gradient = False - - indexer_backend = getattr( - self.config, "csa_indexer_backend", "tilelang" - ) - use_tilelang_indexer = indexer_backend == "tilelang" - use_tilelang_loss_path = ( - use_tilelang_indexer - and self.training - and paddle.is_grad_enabled() - ) - loss_topk_effective = _resolve_csa_indexer_loss_topk_effective( - self.config, self.indexer.index_topk, n_compressed_global - ) - attn_topk_effective = _resolve_csa_indexer_attn_topk_effective( - self.indexer.index_topk, n_compressed_global - ) - - # valid_range for varlen: [b, sq_local, 2] or None - if startend_row_indices is not None: - valid_range_full = get_valid_range( - int(self.compress_ratio), + indexcache_action = self._indexcache_next_c4_action() + if ( + indexcache_action is not None + and indexcache_action[1] == "S" + ): + c4_ordinal, _action, pattern = indexcache_action + compress_topk_idxs = self._indexcache_reuse_topk( b, - sq_global, - startend_row_indices, + sq, + c4_ordinal, + pattern, + indexcache_state=indexcache_state, ) - valid_range = valid_range_full[ - :, position_offset : position_offset + sq, : - ] + indexcache_served_loss_state = ( + self._indexcache_served_distill_state( + query, + compressed_kv_global, + c4_ordinal, + pattern, + loss_mask=loss_mask, + global_valid_count=global_valid_count, + indexcache_state=indexcache_state, + ) + ) + if self._indexcache_has_future_served_layer( + pattern, c4_ordinal + ): + indexcache_state_next = indexcache_state + else: + indexcache_state_next = None + self._indexcache_clear_cached_state() + indexcache_state_updated = True else: - valid_range = None - - q_indexer_bf, k_indexer_global, weights_indexer_bf = ( - self.indexer.forward_before_topk( - x_det, - qr_det, - startend_row_indices=startend_row_indices, - position_offset=position_offset, - cp_group=self.cp_group, + x_det = x.detach() + qr_det = qr.detach() + if self.training: + x_det.stop_gradient = False + qr_det.stop_gradient = False + + indexer_backend = getattr( + self.config, "csa_indexer_backend", "tilelang" ) - ) - - indexer_loss_coeff = getattr( - self.config, "dsa_indexer_loss_coeff", 0.0 - ) - - if use_tilelang_loss_path: # CP training grad-enabled forward with TileLang indexer backend. - # Fused TileLang: single PyLayer produces topk + loss. - # key_comp_mla is 3D [b, n_comp_global, hn] (shared across heads). - key_comp_mla = compressed_kv_global.detach() - ( - indexer_loss, - topk_indices_compressed, - topk_probs, - target, - ) = _compute_fused_csa_indexer_loss_forward( - q_indexer_bf, - weights_indexer_bf, - k_indexer_global, - query.detach(), - key_comp_mla, - valid_range, - int(self.compress_ratio), - int(loss_topk_effective), - float(self.softmax_scale), - float(indexer_loss_coeff), - self.tp_group, - seq_offset=position_offset, - loss_mask=loss_mask, - global_valid_count=global_valid_count, + use_tilelang_indexer = indexer_backend == "tilelang" + use_tilelang_loss_path = ( + use_tilelang_indexer + and self.training + and paddle.is_grad_enabled() ) - tilelang_indexer_loss_state = ( - q_indexer_bf, - weights_indexer_bf, - k_indexer_global, - topk_indices_compressed, - topk_probs, - target, - float(indexer_loss_coeff) - if loss_mask is not None - else float(indexer_loss_coeff) / self.cp_size, - getattr(self.config, "csa_indexer_backend", "tilelang"), - global_valid_count if loss_mask is not None else None, - loss_mask, + loss_topk_effective = ( + _resolve_csa_indexer_loss_topk_effective( + self.config, + self.indexer.index_topk, + n_compressed_global, + ) ) - if ( - indexer_loss_coeff > 0 - ): # CP TileLang training path logs only when indexer loss is enabled. - DSAIndexerLossLoggingHelper.save_loss_to_tracker( - loss=indexer_loss, - layer_number=self.layer_number, - num_layers=self.config.num_hidden_layers, + attn_topk_effective = ( + _resolve_csa_indexer_attn_topk_effective( + self.indexer.index_topk, n_compressed_global ) - # Scale: each rank's loss is mean over sq_local; - # global loss is mean over sq_global = sq_local * cp_size. - if loss_mask is None: - indexer_loss = indexer_loss / self.cp_size - - elif ( - self.training and not use_tilelang_indexer - ): # CP training forward with unfused indexer backend. - # Paddle reference loss path - key_for_loss = ( - compressed_kv_global.detach() - .unsqueeze(2) - .expand([-1, -1, np_heads, -1]) ) - if startend_row_indices is None: - causal_mask = build_causal_mask_cp( - q_positions, - n_compressed_global, - self.compress_ratio, - b, - ) - else: - causal_mask_full = _build_compressed_causal_mask( - self.compress_ratio, + # valid_range for varlen: [b, sq_local, 2] or None + if startend_row_indices is not None: + valid_range_full = get_valid_range( + int(self.compress_ratio), b, sq_global, - n_compressed_global, startend_row_indices, ) - causal_mask = causal_mask_full[ - :, position_offset : position_offset + sq, ... + valid_range = valid_range_full[ + :, position_offset : position_offset + sq, : ] + else: + valid_range = None - weights_for_loss = ( - weights_indexer_bf * self.indexer.softmax_scale - ) - mask_for_loss = causal_mask.unsqueeze(1) - - indexer_loss = FusedDSAIndexerLoss.apply( - q_indexer_bf, - weights_for_loss, - k_indexer_global, - query.detach(), - key_for_loss.detach(), - self.softmax_scale, - min(self.indexer.index_topk, n_compressed_global), - indexer_loss_coeff, - mask_for_loss, - getattr( - self.config, "dsa_indexer_use_sparse_loss", True - ), - self.tp_group, - loss_mask, - global_valid_count, - ) - topk_indices_compressed = ( - FusedDSAIndexerLoss._last_topk_indices + q_indexer_bf, k_indexer_global, weights_indexer_bf = ( + self.indexer.forward_before_topk( + x_det, + qr_det, + startend_row_indices=startend_row_indices, + position_offset=position_offset, + cp_group=self.cp_group, + ) ) + + served_count = None + loss_coeff_override = None + materialize_distill_state = False if ( - indexer_loss_coeff > 0 - ): # CP unfused training path logs only when indexer loss is enabled. - DSAIndexerLossLoggingHelper.save_loss_to_tracker( - loss=indexer_loss, - layer_number=self.layer_number, - num_layers=self.config.num_hidden_layers, + indexcache_action is not None + and self._indexcache_multi_layer_distill_enabled() + ): + c4_ordinal, _action, pattern = indexcache_action + served_count = self._indexcache_served_count( + pattern, c4_ordinal ) - if loss_mask is None: - indexer_loss = indexer_loss / self.cp_size - - elif not use_tilelang_indexer: # CP eval/no-grad forward with unfused backend; only materialize attention top-k. - # Inference-only Paddle topk (use already-gathered global K) - if startend_row_indices is None: - causal_mask = build_causal_mask_cp( - q_positions, - n_compressed_global, - self.compress_ratio, - b, + loss_coeff_override = ( + self._indexcache_scaled_loss_coeff(served_count) ) - else: - causal_mask_full = _build_compressed_causal_mask( - self.compress_ratio, - b, - sq_global, - n_compressed_global, - startend_row_indices, + materialize_distill_state = ( + self._indexcache_has_future_served_layer( + pattern, c4_ordinal + ) ) - causal_mask = causal_mask_full[ - :, position_offset : position_offset + sq, ... - ] - - _, topk_indices_compressed = fused_qk_topk_naive( - q_indexer_bf, - k_indexer_global, - weights_indexer_bf, - attn_topk_effective, - causal_mask, + indexer_loss_coeff = ( + loss_coeff_override + if loss_coeff_override is not None + else getattr(self.config, "dsa_indexer_loss_coeff", 0.0) ) - # TileLang fwd-only topk (no loss, or loss already produced above) - if ( - use_tilelang_indexer and not use_tilelang_loss_path - ): # CP eval/no-grad or recompute first forward with TileLang backend. - from paddlefleet.tilelang_ops import csa_indexer_topk_fwd + if use_tilelang_loss_path: # CP training grad-enabled forward with TileLang indexer backend. + # Fused TileLang: single PyLayer produces topk + loss. + # key_comp_mla is 3D [b, n_comp_global, hn] (shared across heads). + key_comp_mla = compressed_kv_global.detach() + ( + indexer_loss, + topk_indices_compressed, + topk_probs, + target, + ) = _compute_fused_csa_indexer_loss_forward( + q_indexer_bf, + weights_indexer_bf, + k_indexer_global, + query.detach(), + key_comp_mla, + valid_range, + int(self.compress_ratio), + int(loss_topk_effective), + float(self.softmax_scale), + float(indexer_loss_coeff), + self.tp_group, + seq_offset=position_offset, + loss_mask=loss_mask, + global_valid_count=global_valid_count, + ) + tilelang_indexer_loss_state = ( + q_indexer_bf, + weights_indexer_bf, + k_indexer_global, + topk_indices_compressed, + topk_probs, + target, + float(indexer_loss_coeff) + if loss_mask is not None + else float(indexer_loss_coeff) / self.cp_size, + getattr( + self.config, "csa_indexer_backend", "tilelang" + ), + global_valid_count + if loss_mask is not None + else None, + loss_mask, + ) + if ( + indexer_loss_coeff > 0 + ): # CP TileLang training path logs only when indexer loss is enabled. + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_hidden_layers, + ) + # Scale: each rank's loss is mean over sq_local; + # global loss is mean over sq_global = sq_local * cp_size. + if loss_mask is None: + indexer_loss = indexer_loss / self.cp_size + + elif ( + self.training and not use_tilelang_indexer + ): # CP training forward with unfused indexer backend. + # Paddle reference loss path + key_for_loss = ( + compressed_kv_global.detach() + .unsqueeze(2) + .expand([-1, -1, np_heads, -1]) + ) - with paddle.no_grad(): - tl_topk_indices, _ = csa_indexer_topk_fwd( + if startend_row_indices is None: + causal_mask = build_causal_mask_cp( + q_positions, + n_compressed_global, + self.compress_ratio, + b, + ) + else: + causal_mask_full = _build_compressed_causal_mask( + self.compress_ratio, + b, + sq_global, + n_compressed_global, + startend_row_indices, + ) + causal_mask = causal_mask_full[ + :, position_offset : position_offset + sq, ... + ] + + weights_for_loss = ( + weights_indexer_bf * self.indexer.softmax_scale + ) + mask_for_loss = causal_mask.unsqueeze(1) + + indexer_loss = FusedDSAIndexerLoss.apply( + q_indexer_bf, + weights_for_loss, + k_indexer_global, + query.detach(), + key_for_loss.detach(), + self.softmax_scale, + min(self.indexer.index_topk, n_compressed_global), + indexer_loss_coeff, + mask_for_loss, + getattr( + self.config, + "dsa_indexer_use_sparse_loss", + True, + ), + self.tp_group, + loss_mask, + global_valid_count, + ) + topk_indices_compressed = ( + FusedDSAIndexerLoss._last_topk_indices + ) + if ( + indexer_loss_coeff > 0 + ): # CP unfused training path logs only when indexer loss is enabled. + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_hidden_layers, + ) + if loss_mask is None: + indexer_loss = indexer_loss / self.cp_size + + elif not use_tilelang_indexer: # CP eval/no-grad forward with unfused backend; only materialize attention top-k. + # Inference-only Paddle topk (use already-gathered global K) + if startend_row_indices is None: + causal_mask = build_causal_mask_cp( + q_positions, + n_compressed_global, + self.compress_ratio, + b, + ) + else: + causal_mask_full = _build_compressed_causal_mask( + self.compress_ratio, + b, + sq_global, + n_compressed_global, + startend_row_indices, + ) + causal_mask = causal_mask_full[ + :, position_offset : position_offset + sq, ... + ] + + _, topk_indices_compressed = fused_qk_topk_naive( q_indexer_bf, k_indexer_global, weights_indexer_bf, - ratio=self.compress_ratio, - topk_effective=attn_topk_effective, - seq_offset=position_offset, - valid_range=valid_range, + attn_topk_effective, + causal_mask, ) - topk_indices_compressed = tl_topk_indices - if ( - topk_indices_compressed.shape[-1] > attn_topk_effective - ): # CP loss path may return wider top-k than attention consumes. - topk_indices_compressed = topk_indices_compressed[ - ..., :attn_topk_effective - ].contiguous() + # TileLang fwd-only topk (no loss, or loss already produced above) + if ( + use_tilelang_indexer and not use_tilelang_loss_path + ): # CP eval/no-grad or recompute first forward with TileLang backend. + from paddlefleet.tilelang_ops import csa_indexer_topk_fwd + + with paddle.no_grad(): + tl_topk_effective = ( + loss_topk_effective + if materialize_distill_state + else attn_topk_effective + ) + tl_topk_indices, tl_topk_probs = csa_indexer_topk_fwd( + q_indexer_bf, + k_indexer_global, + weights_indexer_bf, + ratio=self.compress_ratio, + topk_effective=tl_topk_effective, + seq_offset=position_offset, + valid_range=valid_range, + ) + topk_indices_compressed = tl_topk_indices + if materialize_distill_state: + tilelang_indexer_loss_state = ( + q_indexer_bf, + weights_indexer_bf, + k_indexer_global, + topk_indices_compressed, + tl_topk_probs, + None, + float(indexer_loss_coeff) + if loss_mask is not None + else float(indexer_loss_coeff) / self.cp_size, + getattr( + self.config, + "csa_indexer_backend", + "tilelang", + ), + global_valid_count + if loss_mask is not None + else None, + loss_mask, + ) - compress_topk_idxs = map_compressed_topk_to_kv_full_cp( - topk_indices_compressed, - q_positions, - self.compress_ratio, - offset, - ) + if ( + topk_indices_compressed.shape[-1] + > attn_topk_effective + ): # CP loss path may return wider top-k than attention consumes. + topk_indices_compressed = topk_indices_compressed[ + ..., :attn_topk_effective + ].contiguous() + + compress_topk_idxs = map_compressed_topk_to_kv_full_cp( + topk_indices_compressed, + q_positions, + self.compress_ratio, + offset, + ) + if indexcache_action is not None: + c4_ordinal, _action, pattern = indexcache_action + effective_loss_scale = ( + tilelang_indexer_loss_state[6] + if tilelang_indexer_loss_state is not None + else loss_coeff_override + ) + produced_state = self._indexcache_cache_topk( + compress_topk_idxs, + c4_ordinal, + pattern, + tilelang_indexer_loss_state, + served_count, + effective_loss_scale, + ) + if self._indexcache_has_future_served_layer( + pattern, c4_ordinal + ): + indexcache_state_next = produced_state + else: + indexcache_state_next = None + self._indexcache_clear_cached_state() + indexcache_state_updated = True else: # HCA path: attend to all compressed positions if startend_row_indices is None: @@ -2368,13 +3202,29 @@ def _forward_cp( ) # Step 5: Attach indexer loss - if tilelang_indexer_loss_state is not None and self.training: + if ( + tilelang_indexer_loss_state is not None + and self.training + and paddle.is_grad_enabled() + ): output = TileLangCSAIndexerLossAutoScaler.apply( output, *tilelang_indexer_loss_state ) - elif indexer_loss is not None and self.training: + elif ( + indexer_loss is not None + and self.training + and paddle.is_grad_enabled() + ): output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + if indexcache_served_loss_state is not None: + output = TileLangCSAIndexerLossAutoScaler.apply( + output, *indexcache_served_loss_state + ) + if indexcache_state_next is not None: + return output, indexcache_state_next + if indexcache_state_updated: + return output, None return output def compressed_sparse_attn( diff --git a/src/paddlefleet/transformer/dsv4_hybrid_attention.py b/src/paddlefleet/transformer/dsv4_hybrid_attention.py index bdab04fdf5..0fc44f4411 100644 --- a/src/paddlefleet/transformer/dsv4_hybrid_attention.py +++ b/src/paddlefleet/transformer/dsv4_hybrid_attention.py @@ -296,8 +296,8 @@ def forward( ) # Core attention (CompressedSparseAttention) - input_ids = kwargs.get("input_ids", None) - core_attn_out = self.core_attention( + indexcache_state = kwargs.get("indexcache_state", None) + core_attn_result = self.core_attention( query, key, value, @@ -306,7 +306,13 @@ def forward( x=hidden_states, qr=q_compressed, input_ids=kwargs.get("input_ids", None), + indexcache_state=indexcache_state, ) + core_attn_returns_indexcache_state = isinstance(core_attn_result, tuple) + if isinstance(core_attn_result, tuple): + core_attn_out, indexcache_state = core_attn_result + else: + core_attn_out = core_attn_result # core_attn_out: [b, sq, np * v_head_dim] # Inverse RoPE on last qk_pos_emb_head_dim of each head @@ -384,6 +390,8 @@ def forward( # Output projection output, bias = self.o_proj(core_attn_out) + if indexcache_state is not None or core_attn_returns_indexcache_state: + return output, bias, indexcache_state return output, bias def get_query_key_value_tensors( diff --git a/src/paddlefleet/transformer/indexcache_state.py b/src/paddlefleet/transformer/indexcache_state.py new file mode 100644 index 0000000000..ae26d2a89d --- /dev/null +++ b/src/paddlefleet/transformer/indexcache_state.py @@ -0,0 +1,132 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +from __future__ import annotations + +from collections.abc import Iterable + +import paddle + + +INDEXCACHE_STATE_KIND_NONE = "none" +INDEXCACHE_STATE_KIND_TOPK_ONLY = "topk_only" +INDEXCACHE_STATE_KIND_DISTILL = "distill" +INDEXCACHE_STATE_KIND_INVALID = "invalid" + +INDEXCACHE_TOPK_ONLY_STATE_LEN = 3 +INDEXCACHE_DISTILL_STATE_LEN = 8 +INDEXCACHE_RECOMPUTE_STATE_MAX_LEN = INDEXCACHE_DISTILL_STATE_LEN + +INDEXCACHE_STATE_TOPK_IDXS = 0 +INDEXCACHE_TOPK_ONLY_STATE_PRODUCER_LAYER = 1 +INDEXCACHE_TOPK_ONLY_STATE_SERVED_COUNT = 2 + +INDEXCACHE_DISTILL_STATE_Q = 1 +INDEXCACHE_DISTILL_STATE_WEIGHTS = 2 +INDEXCACHE_DISTILL_STATE_K = 3 +INDEXCACHE_DISTILL_STATE_TOPK_INDICES = 4 +INDEXCACHE_DISTILL_STATE_TOPK_PROBS = 5 +INDEXCACHE_DISTILL_STATE_PRODUCER_LAYER = 6 +INDEXCACHE_DISTILL_STATE_SERVED_COUNT = 7 + +INDEXCACHE_DISTILL_GRAD_INDICES = ( + INDEXCACHE_DISTILL_STATE_Q, + INDEXCACHE_DISTILL_STATE_WEIGHTS, + INDEXCACHE_DISTILL_STATE_K, +) + + +def state_kind(indexcache_state: tuple | list | None) -> str: + if not indexcache_state: + return INDEXCACHE_STATE_KIND_NONE + state_len = len(indexcache_state) + if state_len == INDEXCACHE_TOPK_ONLY_STATE_LEN: + return INDEXCACHE_STATE_KIND_TOPK_ONLY + if state_len == INDEXCACHE_DISTILL_STATE_LEN: + return INDEXCACHE_STATE_KIND_DISTILL + return INDEXCACHE_STATE_KIND_INVALID + + +def is_valid_state(indexcache_state: tuple | list | None) -> bool: + return state_kind(indexcache_state) in ( + INDEXCACHE_STATE_KIND_NONE, + INDEXCACHE_STATE_KIND_TOPK_ONLY, + INDEXCACHE_STATE_KIND_DISTILL, + ) + + +def _as_state_tuple(indexcache_state: tuple | list | None) -> tuple | None: + if indexcache_state is None: + return None + return tuple(indexcache_state) + + +def apply_stop_gradient_mask( + indexcache_state: tuple | list | None, +) -> tuple | None: + indexcache_state = _as_state_tuple(indexcache_state) + kind = state_kind(indexcache_state) + if kind == INDEXCACHE_STATE_KIND_NONE: + return None + if kind == INDEXCACHE_STATE_KIND_INVALID: + raise ValueError( + "IndexCache state must be either topk-only " + f"({INDEXCACHE_TOPK_ONLY_STATE_LEN} tensors) or distill " + f"({INDEXCACHE_DISTILL_STATE_LEN} tensors), got " + f"len={len(indexcache_state)}." + ) + + for idx, tensor in enumerate(indexcache_state): + if not isinstance(tensor, paddle.Tensor): + raise TypeError( + "IndexCache state entries must be Paddle tensors, got " + f"type={type(tensor).__name__} at index={idx}." + ) + tensor.stop_gradient = not ( + kind == INDEXCACHE_STATE_KIND_DISTILL + and idx in INDEXCACHE_DISTILL_GRAD_INDICES + ) + return indexcache_state + + +def detach_stop_gradient_tensor(value: paddle.Tensor) -> paddle.Tensor: + value = value.detach() + value.stop_gradient = True + return value + + +def clone_state_outputs(indexcache_state: tuple | list | None) -> tuple | None: + masked = apply_stop_gradient_mask(indexcache_state) + if masked is None: + return None + cloned = tuple(tensor.clone() for tensor in masked) + return apply_stop_gradient_mask(cloned) + + +def flatten_state(indexcache_state: tuple | list | None) -> tuple: + masked = apply_stop_gradient_mask(indexcache_state) + if masked is None: + return () + return masked + + +def state_to_slots(indexcache_state: tuple | list | None) -> tuple: + slots = list(flatten_state(indexcache_state)) + if len(slots) > INDEXCACHE_RECOMPUTE_STATE_MAX_LEN: + raise ValueError( + "IndexCache recompute state has too many slots: " + f"{len(slots)} > {INDEXCACHE_RECOMPUTE_STATE_MAX_LEN}." + ) + slots.extend([None] * (INDEXCACHE_RECOMPUTE_STATE_MAX_LEN - len(slots))) + return tuple(slots) + + +def state_from_slots(slots: Iterable) -> tuple | None: + state = list(slots) + while state and state[-1] is None: + state.pop() + if not state: + return None + return apply_stop_gradient_mask(tuple(state)) diff --git a/src/paddlefleet/transformer/moe/fp8_utils.py b/src/paddlefleet/transformer/moe/fp8_utils.py index c1d2a145cb..499db751fb 100644 --- a/src/paddlefleet/transformer/moe/fp8_utils.py +++ b/src/paddlefleet/transformer/moe/fp8_utils.py @@ -1902,6 +1902,22 @@ def bf16_weight_grad(self, dy, x, weights, p2p_overlap=False): """ BF16 GEMM for weight grad """ + def _get_stop_gradient(w): + if isinstance(w, (list, tuple)): + return all(getattr(item, "stop_gradient", True) for item in w) + parent = getattr(w, "_parent", None) + if parent is not None: + for attr in ("weight1", "weight2"): + parent_weight = getattr(parent, attr, None) + if parent_weight is not None and hasattr( + parent_weight, "stop_gradient" + ): + return bool(parent_weight.stop_gradient) + return bool(getattr(w, "stop_gradient", True)) + + if _get_stop_gradient(weights): + return + if x is None: if self.dequant_input: x = paddle.incubate.nn.functional.fused_act_dequant( diff --git a/src/paddlefleet/transformer/moe/moe_router.py b/src/paddlefleet/transformer/moe/moe_router.py index cd9edadef1..237989d0b8 100644 --- a/src/paddlefleet/transformer/moe/moe_router.py +++ b/src/paddlefleet/transformer/moe/moe_router.py @@ -974,7 +974,7 @@ def _setup_hash_layer(self, layer_number, is_mtp_layer: bool = False): # inference checkpoint; no public initialization recipe is documented. # Round-robin is used here only as a placeholder so the layer is # runnable from scratch. - ids = paddle.arange(vocab_size, dtype=paddle.int64) + ids = paddle.arange(vocab_size, dtype=paddle.int32) tid2eid = paddle.stack( [ (ids + k) % self.num_experts diff --git a/src/paddlefleet/transformer/transformer_config.py b/src/paddlefleet/transformer/transformer_config.py index 07524e6a65..3699f76deb 100644 --- a/src/paddlefleet/transformer/transformer_config.py +++ b/src/paddlefleet/transformer/transformer_config.py @@ -874,6 +874,20 @@ class TransformerConfig(ModelParallelConfig): kernel. """ + index_topk_pattern: str | None = None + """Optional IndexCache training pattern over ratio=4 CSA indexer layers. + + Each character corresponds to one ratio=4 layer in execution order: + - "F": run the learned indexer and cache its top-k indices + - "S": skip the local indexer and reuse the previous cached top-k indices + """ + + indexcache_multi_layer_distill: bool = False + """Train retained IndexCache indexers with targets from all served layers. + + CP training is supported only for no-MTP TileLang CSA. + """ + use_fast_hadamard: bool = False """Use Tridao's fast Hadamard transform for DSv4 rotate activation function.""" @@ -930,6 +944,8 @@ class TransformerConfig(ModelParallelConfig): "csa_dense_mode": "csa_dense_mode", "csa_indexer_backend": "csa_indexer_backend", "csa_sparse_attn_backend": "csa_sparse_attn_backend", + "index_topk_pattern": "index_topk_pattern", + "indexcache_multi_layer_distill": "indexcache_multi_layer_distill", "o_groups": "o_groups", "o_lora_rank": "o_lora_rank", "qk_pos_emb_head_dim": "qk_pos_emb_head_dim", @@ -1189,6 +1205,85 @@ def __post_init__(self): f"Must be one of {valid_ratios}." ) + if self.index_topk_pattern is not None: + pattern = str(self.index_topk_pattern).strip().upper() + if not pattern: + self.index_topk_pattern = None + else: + invalid_chars = sorted(set(pattern) - {"F", "S"}) + if invalid_chars: + raise ValueError( + "index_topk_pattern may only contain 'F' and 'S', " + f"got invalid chars: {invalid_chars}." + ) + if pattern[0] != "F": + raise ValueError( + "index_topk_pattern must start with 'F' so there " + "is a producer before any reused top-k indices." + ) + if self.csa_dense_mode: + raise ValueError( + "index_topk_pattern requires csa_dense_mode=False." + ) + c4_layer_count = sum( + 1 for ratio in self.csa_compress_ratios if ratio == 4 + ) + if len(pattern) != c4_layer_count: + raise ValueError( + "index_topk_pattern length must equal the number " + f"of ratio=4 CSA layers ({c4_layer_count}), got " + f"{len(pattern)}." + ) + self.index_topk_pattern = pattern + + if self.indexcache_multi_layer_distill: + if not self.index_topk_pattern: + raise ValueError( + "indexcache_multi_layer_distill=True requires a " + "non-empty index_topk_pattern." + ) + if self.context_parallel_size > 1 and ( + self.num_nextn_predict_layers is not None + and self.num_nextn_predict_layers > 0 + ): + raise NotImplementedError( + "indexcache_multi_layer_distill currently supports " + "CP training only when num_nextn_predict_layers=0." + ) + if self.csa_indexer_backend != "tilelang": + raise NotImplementedError( + "indexcache_multi_layer_distill currently supports " + "only csa_indexer_backend='tilelang'." + ) + if self.csa_sparse_attn_backend != "tilelang": + raise NotImplementedError( + "indexcache_multi_layer_distill currently supports " + "only csa_sparse_attn_backend='tilelang'." + ) + + if self.index_topk_pattern and self.recompute_granularity: + if not ( + self.recompute_granularity == "full" + and self.recompute_method == "uniform" + and self.recompute_num_layers == 1 + ): + raise NotImplementedError( + "IndexCache recompute currently supports only " + "recompute_granularity='full', " + "recompute_method='uniform', and " + "recompute_num_layers=1." + ) + if self.csa_indexer_backend != "tilelang": + raise NotImplementedError( + "IndexCache recompute currently supports only " + "csa_indexer_backend='tilelang'." + ) + if self.csa_sparse_attn_backend != "tilelang": + raise NotImplementedError( + "IndexCache recompute currently supports only " + "csa_sparse_attn_backend='tilelang'." + ) + if ( getattr(self, "csa_tilelang_enable_sparse_attn", None) is not None diff --git a/src/paddlefleet/transformer/transformer_layer.py b/src/paddlefleet/transformer/transformer_layer.py index 472c519e5f..45365f1d2b 100644 --- a/src/paddlefleet/transformer/transformer_layer.py +++ b/src/paddlefleet/transformer/transformer_layer.py @@ -43,6 +43,18 @@ ) from paddlefleet.transformer.dsv4_hybrid_attention import DSv4HybridAttention from paddlefleet.transformer.identity_op import IdentityFuncOp, IdentityOp +from paddlefleet.transformer.indexcache_state import ( + INDEXCACHE_DISTILL_STATE_LEN, + INDEXCACHE_RECOMPUTE_STATE_MAX_LEN, + INDEXCACHE_STATE_KIND_DISTILL, + INDEXCACHE_STATE_KIND_TOPK_ONLY, + INDEXCACHE_TOPK_ONLY_STATE_LEN, + apply_stop_gradient_mask, + clone_state_outputs, + state_from_slots, + state_kind, + state_to_slots, +) from paddlefleet.transformer.mlp import MLP from paddlefleet.transformer.moe.moe_layer import MoELayer from paddlefleet.transformer.utils import profile @@ -94,6 +106,100 @@ def tensors_clone(outputs): ) +def _is_indexcache_recompute_state(value): + if not ( + isinstance(value, (tuple, list)) + and len(value) + in (INDEXCACHE_TOPK_ONLY_STATE_LEN, INDEXCACHE_DISTILL_STATE_LEN) + and all(isinstance(item, paddle.Tensor) for item in value) + ): + return False + return state_kind(value) in ( + INDEXCACHE_STATE_KIND_TOPK_ONLY, + INDEXCACHE_STATE_KIND_DISTILL, + ) + + +def _mark_indexcache_recompute_state_stop_gradient(value): + if not _is_indexcache_recompute_state(value): + return value + return apply_stop_gradient_mask(value) + + +def _ensure_recompute_non_leaf_tensor(value): + if ( + isinstance(value, paddle.Tensor) + and not value.stop_gradient + and getattr(value, "is_leaf", False) + ): + return paddle.scale(value, scale=1.0, bias=0.0) + return value + + +def _describe_indexcache_recompute_input(value): + if value is None: + return None + if isinstance(value, paddle.Tensor): + return { + "type": "Tensor", + "shape": list(value.shape), + "dtype": str(value.dtype), + "stop_gradient": bool(value.stop_gradient), + "is_leaf": bool(getattr(value, "is_leaf", False)), + } + if isinstance(value, (tuple, list)): + return [ + _describe_indexcache_recompute_input(item) + for item in value + if item is not None + ] + return {"type": type(value).__name__} + + +def _clone_indexcache_recompute_state_outputs(value): + return clone_state_outputs(value) + + +def _flatten_indexcache_recompute_outputs(outputs): + if not isinstance(outputs, tuple) or len(outputs) <= 2: + return outputs + + output, context, indexcache_state = outputs[0], outputs[1], outputs[2] + if not _is_indexcache_recompute_state(indexcache_state): + return outputs + + indexcache_state = _clone_indexcache_recompute_state_outputs(indexcache_state) + + if context is None: + return (output, *indexcache_state) + return (output, context, *indexcache_state) + + +def _unpack_flattened_indexcache_recompute_outputs(outputs): + if not isinstance(outputs, tuple): + return None + + if ( + len(outputs) in ( + 1 + INDEXCACHE_TOPK_ONLY_STATE_LEN, + 1 + INDEXCACHE_DISTILL_STATE_LEN, + ) + and all(isinstance(item, paddle.Tensor) for item in outputs[1:]) + ): + indexcache_state = apply_stop_gradient_mask(outputs[1:]) + return outputs[0], None, indexcache_state + + if ( + len(outputs) + in (2 + INDEXCACHE_TOPK_ONLY_STATE_LEN, 2 + INDEXCACHE_DISTILL_STATE_LEN) + and all(isinstance(item, paddle.Tensor) for item in outputs[2:]) + ): + indexcache_state = apply_stop_gradient_mask(outputs[2:]) + return outputs[0], outputs[1], indexcache_state + + return None + + @dataclass class TransformerLayerSublayersSpec: """ @@ -595,8 +701,122 @@ def forward( attention_bias = dict_args.get("attention_bias", None) packed_seq_params = dict_args.get("packed_seq_params", None) input_ids = dict_args.get("input_ids", None) + indexcache_state = dict_args.get("indexcache_state", None) + indexcache_state = ( + apply_stop_gradient_mask(indexcache_state) + if indexcache_state is not None + else None + ) + indexcache_state_slots = state_to_slots(indexcache_state) + assert ( + len(indexcache_state_slots) == INDEXCACHE_RECOMPUTE_STATE_MAX_LEN + ) + hidden_states = _ensure_recompute_non_leaf_tensor(hidden_states) + attention_mask = _ensure_recompute_non_leaf_tensor(attention_mask) + context = _ensure_recompute_non_leaf_tensor(context) + context_mask = _ensure_recompute_non_leaf_tensor(context_mask) + attn_mask_startend_row_indices = _ensure_recompute_non_leaf_tensor( + attn_mask_startend_row_indices + ) + rotary_pos_emb = _ensure_recompute_non_leaf_tensor(rotary_pos_emb) + rotary_pos_cos = _ensure_recompute_non_leaf_tensor(rotary_pos_cos) + rotary_pos_sin = _ensure_recompute_non_leaf_tensor(rotary_pos_sin) + swa_rotary_pos_emb = _ensure_recompute_non_leaf_tensor( + swa_rotary_pos_emb + ) + swa_rotary_pos_cos = _ensure_recompute_non_leaf_tensor( + swa_rotary_pos_cos + ) + swa_rotary_pos_sin = _ensure_recompute_non_leaf_tensor( + swa_rotary_pos_sin + ) + position_ids = _ensure_recompute_non_leaf_tensor(position_ids) + attention_bias = _ensure_recompute_non_leaf_tensor(attention_bias) + input_ids = _ensure_recompute_non_leaf_tensor(input_ids) + + if os.environ.get("INDEXCACHE_TRAIN_DEBUG", "0") == "1": + print( + "[INDEXCACHE_RECOMPUTE_INPUT] " + f"layer_number={self.layer_number} " + f"hidden_states={_describe_indexcache_recompute_input(hidden_states)} " + f"attention_mask={_describe_indexcache_recompute_input(attention_mask)} " + f"attn_mask_startend_row_indices={_describe_indexcache_recompute_input(attn_mask_startend_row_indices)} " + f"context={_describe_indexcache_recompute_input(context)} " + f"context_mask={_describe_indexcache_recompute_input(context_mask)} " + f"rotary_pos_emb={_describe_indexcache_recompute_input(rotary_pos_emb)} " + f"rotary_pos_cos={_describe_indexcache_recompute_input(rotary_pos_cos)} " + f"rotary_pos_sin={_describe_indexcache_recompute_input(rotary_pos_sin)} " + f"swa_rotary_pos_emb={_describe_indexcache_recompute_input(swa_rotary_pos_emb)} " + f"swa_rotary_pos_cos={_describe_indexcache_recompute_input(swa_rotary_pos_cos)} " + f"swa_rotary_pos_sin={_describe_indexcache_recompute_input(swa_rotary_pos_sin)} " + f"position_ids={_describe_indexcache_recompute_input(position_ids)} " + f"attention_bias={_describe_indexcache_recompute_input(attention_bias)} " + f"packed_seq_params={_describe_indexcache_recompute_input(packed_seq_params)} " + f"input_ids={_describe_indexcache_recompute_input(input_ids)} " + f"indexcache_state={_describe_indexcache_recompute_input(indexcache_state)}", + flush=True, + ) + + def _forward_impl_for_recompute( + hidden_states, + attention_mask=None, + attn_mask_startend_row_indices=None, + context=None, + context_mask=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + swa_rotary_pos_emb=None, + swa_rotary_pos_cos=None, + swa_rotary_pos_sin=None, + position_ids=None, + attention_bias=None, + packed_seq_params=None, + input_ids=None, + indexcache_state_0=None, + indexcache_state_1=None, + indexcache_state_2=None, + indexcache_state_3=None, + indexcache_state_4=None, + indexcache_state_5=None, + indexcache_state_6=None, + indexcache_state_7=None, + ): + indexcache_state = state_from_slots( + ( + indexcache_state_0, + indexcache_state_1, + indexcache_state_2, + indexcache_state_3, + indexcache_state_4, + indexcache_state_5, + indexcache_state_6, + indexcache_state_7, + ) + ) + return _flatten_indexcache_recompute_outputs( + self._forward_impl( + hidden_states=hidden_states, + attention_mask=attention_mask, + attn_mask_startend_row_indices=attn_mask_startend_row_indices, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + swa_rotary_pos_emb=swa_rotary_pos_emb, + swa_rotary_pos_cos=swa_rotary_pos_cos, + swa_rotary_pos_sin=swa_rotary_pos_sin, + position_ids=position_ids, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + input_ids=input_ids, + indexcache_state=indexcache_state, + ) + ) + outputs = recompute( - self._forward_impl, + _forward_impl_for_recompute, hidden_states=hidden_states, attention_mask=attention_mask, attn_mask_startend_row_indices=attn_mask_startend_row_indices.clone() # Clone is necessary! @@ -628,12 +848,30 @@ def forward( attention_bias=attention_bias, packed_seq_params=packed_seq_params, input_ids=input_ids, + indexcache_state_0=indexcache_state_slots[0], + indexcache_state_1=indexcache_state_slots[1], + indexcache_state_2=indexcache_state_slots[2], + indexcache_state_3=indexcache_state_slots[3], + indexcache_state_4=indexcache_state_slots[4], + indexcache_state_5=indexcache_state_slots[5], + indexcache_state_6=indexcache_state_slots[6], + indexcache_state_7=indexcache_state_slots[7], ) else: outputs = self._forward_impl(**dict_args) - if isinstance(outputs, tuple): + indexcache_state = None + flattened_indexcache_outputs = ( + _unpack_flattened_indexcache_recompute_outputs(outputs) + if self.full_recompute + else None + ) + if flattened_indexcache_outputs is not None: + output, context, indexcache_state = flattened_indexcache_outputs + elif isinstance(outputs, tuple): output, context = outputs[0], outputs[1] + if len(outputs) > 2: + indexcache_state = outputs[2] else: output, context = outputs, None @@ -694,7 +932,44 @@ def forward( # dict_args unchanged and will be consumed by MTP layer directly if context is not None: rst["context"] = context + if indexcache_state is not None: + rst["indexcache_state"] = indexcache_state + else: + dict_args.pop("indexcache_state", None) rst = {**dict_args, **rst} + if ( + os.environ.get("INDEXCACHE_TRAIN_DEBUG", "0") == "1" + and getattr(self.config, "index_topk_pattern", None) + ): + state = rst.get("indexcache_state", None) + if isinstance(state, (tuple, list)): + state_len = len(state) + state_shapes = [ + list(item.shape) + if isinstance(item, paddle.Tensor) + else type(item).__name__ + for item in state + ] + state_stop_gradients = [ + bool(item.stop_gradient) + if isinstance(item, paddle.Tensor) + else None + for item in state + ] + else: + state_len = 0 + state_shapes = None + state_stop_gradients = None + print( + "[INDEXCACHE_TRAIN_FLOW] " + f"layer={self.layer_number} " + f"full_recompute={self.full_recompute} " + f"flattened={flattened_indexcache_outputs is not None} " + f"state_len={state_len} state_shapes={state_shapes} " + f"state_stop_gradients={state_stop_gradients} " + f"keys={list(rst.keys())}", + flush=True, + ) return rst def _forward_impl( @@ -714,8 +989,28 @@ def _forward_impl( attention_bias: Tensor | None = None, packed_seq_params: PackedSeqParams | None = None, input_ids: Tensor | None = None, + indexcache_state: tuple | list | None = None, **kwargs, ): + if indexcache_state is None: + indexcache_state = kwargs.get("indexcache_state", None) + if _is_indexcache_recompute_state(indexcache_state): + indexcache_state = _mark_indexcache_recompute_state_stop_gradient( + indexcache_state + ) + kwargs["indexcache_state"] = indexcache_state + elif indexcache_state is not None and "indexcache_state" not in kwargs: + kwargs["indexcache_state"] = indexcache_state + + def unpack_attention_outputs(attention_outputs): + if isinstance(attention_outputs, tuple) and len(attention_outputs) == 3: + return ( + attention_outputs[0], + attention_outputs[1], + attention_outputs[2], + ) + return attention_outputs[0], attention_outputs[1], indexcache_state + def need_do_attention(): # need_do_prefill = forward_meta.max_len_tensor_cpu[1] > 0 # need_do_decode = forward_meta.max_len_tensor_cpu[2] > 0 @@ -760,25 +1055,29 @@ def need_do_attention(): # Self-attention (skip internal bda residual) with profile("attn"): if need_do_attention(): - hidden_states, context = self._forward_attention( - hidden_states=hidden_states, - attention_mask=attention_mask, - attn_mask_startend_row_indices=attn_mask_startend_row_indices, - context=context, - context_mask=context_mask, - rotary_pos_emb=rotary_pos_emb, - rotary_pos_cos=rotary_pos_cos, - rotary_pos_sin=rotary_pos_sin, - swa_rotary_pos_emb=swa_rotary_pos_emb, - swa_rotary_pos_cos=swa_rotary_pos_cos, - swa_rotary_pos_sin=swa_rotary_pos_sin, - position_ids=position_ids, - attention_bias=attention_bias, - packed_seq_params=packed_seq_params, - block_attention_residuals=True, - in_recompute=self.full_recompute, - input_ids=input_ids, - **kwargs, + hidden_states, context, indexcache_state = ( + unpack_attention_outputs( + self._forward_attention( + hidden_states=hidden_states, + attention_mask=attention_mask, + attn_mask_startend_row_indices=attn_mask_startend_row_indices, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + swa_rotary_pos_emb=swa_rotary_pos_emb, + swa_rotary_pos_cos=swa_rotary_pos_cos, + swa_rotary_pos_sin=swa_rotary_pos_sin, + position_ids=position_ids, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + block_attention_residuals=True, + in_recompute=self.full_recompute, + input_ids=input_ids, + **kwargs, + ) + ) ) # Accumulate attn output into partial_block @@ -812,24 +1111,28 @@ def need_do_attention(): self._log_md5(hidden_states, "input", self.layer_number) with profile("attn"): if need_do_attention(): - hidden_states, context = self._forward_attention( - hidden_states=hidden_states, - attention_mask=attention_mask, - attn_mask_startend_row_indices=attn_mask_startend_row_indices, - context=context, - context_mask=context_mask, - rotary_pos_emb=rotary_pos_emb, - rotary_pos_cos=rotary_pos_cos, - rotary_pos_sin=rotary_pos_sin, - swa_rotary_pos_emb=swa_rotary_pos_emb, - swa_rotary_pos_cos=swa_rotary_pos_cos, - swa_rotary_pos_sin=swa_rotary_pos_sin, - position_ids=position_ids, - attention_bias=attention_bias, - packed_seq_params=packed_seq_params, - in_recompute=self.full_recompute, - input_ids=input_ids, - **kwargs, + hidden_states, context, indexcache_state = ( + unpack_attention_outputs( + self._forward_attention( + hidden_states=hidden_states, + attention_mask=attention_mask, + attn_mask_startend_row_indices=attn_mask_startend_row_indices, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + swa_rotary_pos_emb=swa_rotary_pos_emb, + swa_rotary_pos_cos=swa_rotary_pos_cos, + swa_rotary_pos_sin=swa_rotary_pos_sin, + position_ids=position_ids, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + in_recompute=self.full_recompute, + input_ids=input_ids, + **kwargs, + ) + ) ) self._log_md5( hidden_states, "post_attn_residual", self.layer_number @@ -837,6 +1140,8 @@ def need_do_attention(): with profile(timer_name): output = self._forward_mlp(hidden_states, input_ids=input_ids) self._log_md5(output, "layer_output", self.layer_number) + if indexcache_state is not None: + return output, context, indexcache_state if context is not None: return output, context return output @@ -911,6 +1216,11 @@ def _forward_attention( self.self_attn, DSv4HybridAttention ): extra_kwargs["input_ids"] = input_ids + indexcache_state = kwargs.get("indexcache_state", None) + if indexcache_state is not None and isinstance( + self.self_attn, DSv4HybridAttention + ): + extra_kwargs["indexcache_state"] = indexcache_state if rope_freqs_cis is not None: attention_output_with_bias = self.self_attn( @@ -948,6 +1258,13 @@ def _forward_attention( **extra_kwargs, ) + if ( + isinstance(attention_output_with_bias, tuple) + and len(attention_output_with_bias) == 3 + ): + indexcache_state = attention_output_with_bias[2] + attention_output_with_bias = attention_output_with_bias[:2] + with paddle.enable_grad(): if block_attention_residuals: attn_out, attn_bias = attention_output_with_bias @@ -997,6 +1314,8 @@ def _forward_attention( if is_first_fwd: hidden_states.stop_gradient = False + if indexcache_state is not None: + return hidden_states, context, indexcache_state return hidden_states, context def _forward_mlp( @@ -1217,6 +1536,11 @@ def _forward_attention( self.self_attn, DSv4HybridAttention ): extra_kwargs["input_ids"] = kwargs["input_ids"] + indexcache_state = kwargs.get("indexcache_state", None) + if indexcache_state is not None and isinstance( + self.self_attn, DSv4HybridAttention + ): + extra_kwargs["indexcache_state"] = indexcache_state if rope_freqs_cis is not None: attention_output_with_bias = self.self_attn( @@ -1245,6 +1569,13 @@ def _forward_attention( **extra_kwargs, ) + if ( + isinstance(attention_output_with_bias, tuple) + and len(attention_output_with_bias) == 3 + ): + indexcache_state = attention_output_with_bias[2] + attention_output_with_bias = attention_output_with_bias[:2] + # mHC: fused H_res + H_post + bias-dropout-add hidden_states = ( self.self_attention_hyper_connection.fused_h_res_h_post_bda( @@ -1284,6 +1615,8 @@ def _forward_attention( if is_first_fwd: hidden_states.stop_gradient = False + if indexcache_state is not None: + return hidden_states, context, indexcache_state return hidden_states, context def _forward_mlp( diff --git a/tests/single_card_tests/ai_edited_test/moe/test_ai_moe_router.py b/tests/single_card_tests/ai_edited_test/moe/test_ai_moe_router.py index 4ccf205fc9..ce4154812b 100644 --- a/tests/single_card_tests/ai_edited_test/moe/test_ai_moe_router.py +++ b/tests/single_card_tests/ai_edited_test/moe/test_ai_moe_router.py @@ -548,6 +548,7 @@ def test_setup_registers_tid2eid(self): self.assertTrue(router.is_hash_layer) self.assertIsNotNone(router.tid2eid) self.assertEqual(list(router.tid2eid.shape), [16, 2]) + self.assertEqual(router.tid2eid.dtype, paddle.int32) # Round-robin placeholder: tid2eid[i, k] = (i + k) % num_experts. ids = np.arange(16) expected = np.stack([(ids + k) % 4 for k in range(2)], axis=1) diff --git a/tests/single_card_tests/test_indexcache_pp_runtime_patch.py b/tests/single_card_tests/test_indexcache_pp_runtime_patch.py new file mode 100644 index 0000000000..c4a5ed00fe --- /dev/null +++ b/tests/single_card_tests/test_indexcache_pp_runtime_patch.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +import unittest + +import paddle +from paddle.distributed.fleet.meta_parallel.pp_utils.forward_backward_overlap_utils import ( + ScheduleNode, +) +from paddle.distributed.fleet.meta_parallel.pp_utils.p2p_communication import ( + SendRecvMeta, +) + +import paddlefleet # noqa: F401 - installs the IndexCache pipeline patch +from paddlefleet.indexcache_pp_runtime_patch import ( + _dict_to_tuple_helper, + _get_pipeline_key, + _normalize_pipeline_input_gradients, + _tuple_to_dict_helper, +) + + +def _make_distill_state(offset): + return ( + paddle.to_tensor([offset], dtype="int64"), + paddle.to_tensor([1.0 + offset], stop_gradient=False), + paddle.to_tensor([2.0 + offset], stop_gradient=False), + paddle.to_tensor([3.0 + offset], stop_gradient=False), + paddle.to_tensor([offset], dtype="int64"), + paddle.to_tensor([0.5 + offset]), + paddle.to_tensor([offset], dtype="int64"), + paddle.to_tensor([0], dtype="int64"), + ) + + +class TestIndexCachePipelineBackward(unittest.TestCase): + def test_replaced_distill_state_uses_sendable_zero_gradients(self): + old_state = _make_distill_state(0) + new_state = _make_distill_state(10) + node = ScheduleNode( + lambda inputs: { + "hidden": inputs["hidden"] * 2, + "indexcache_state": new_state, + }, + name="replace_indexcache_producer", + ) + + outputs = node.forward( + { + "hidden": paddle.to_tensor([4.0], stop_gradient=False), + "indexcache_state": old_state, + } + ) + output_grads = tuple( + paddle.ones_like(tensor) + for tensor in _dict_to_tuple_helper(outputs) + if not tensor.stop_gradient + ) + input_grads = node.backward(output_grads) + + self.assertEqual(len(input_grads), 4) + self.assertTrue(all(isinstance(grad, paddle.Tensor) for grad in input_grads)) + self.assertTrue(all(not grad.stop_gradient for grad in input_grads)) + self.assertEqual(input_grads[0].item(), 2.0) + for grad in input_grads[1:]: + self.assertEqual(grad.item(), 0.0) + + meta = SendRecvMeta() + meta.set_send_message(input_grads) + self.assertEqual(len(meta.send_shape_message), 4) + + def test_missing_non_indexcache_gradient_fails_fast(self): + node = ScheduleNode( + lambda inputs: {"hidden": inputs["hidden"] * 2}, + name="drop_regular_pipeline_input", + ) + outputs = node.forward( + { + "hidden": paddle.to_tensor([4.0], stop_gradient=False), + "unused": paddle.to_tensor([5.0], stop_gradient=False), + } + ) + output_grads = tuple( + paddle.ones_like(tensor) + for tensor in _dict_to_tuple_helper(outputs) + if not tensor.stop_gradient + ) + + with self.assertRaisesRegex( + RuntimeError, + "missing a gradient outside IndexCache state", + ): + node.backward(output_grads) + + def test_outer_pipeline_boundary_uses_sendable_zero_gradients(self): + pipeline_inputs = _dict_to_tuple_helper( + { + "hidden_states": paddle.to_tensor( + [4.0], stop_gradient=False + ), + "indexcache_state": _make_distill_state(0), + } + ) + converted_inputs, use_dict = _tuple_to_dict_helper(pipeline_inputs) + + self.assertTrue(use_dict) + self.assertIn("indexcache_state", converted_inputs) + self.assertTrue(all(not hasattr(tensor, "key") for tensor in pipeline_inputs)) + + hidden_grad = paddle.ones_like(pipeline_inputs[0]) + hidden_grad.stop_gradient = False + input_grads = _normalize_pipeline_input_gradients( + pipeline_inputs, + (hidden_grad, None, None, None), + ) + + self.assertEqual(len(input_grads), 4) + self.assertTrue(all(isinstance(grad, paddle.Tensor) for grad in input_grads)) + self.assertTrue(all(not grad.stop_gradient for grad in input_grads)) + self.assertEqual(input_grads[0].item(), 1.0) + for grad in input_grads[1:]: + self.assertEqual(grad.item(), 0.0) + + meta = SendRecvMeta() + meta.set_send_message(input_grads) + self.assertEqual(len(meta.send_shape_message), 4) + + def test_outer_pipeline_boundary_preserves_baseline_none_gradient(self): + hidden = paddle.to_tensor([4.0], stop_gradient=False) + unused = paddle.to_tensor([5.0], stop_gradient=False) + input_grads = (paddle.ones_like(hidden), None) + + normalized = _normalize_pipeline_input_gradients( + (hidden, unused), + input_grads, + ) + + self.assertIs(normalized, input_grads) + + def test_outer_pipeline_boundary_handles_released_state(self): + pipeline_inputs = _dict_to_tuple_helper( + { + "hidden_states": paddle.to_tensor( + [4.0], stop_gradient=False + ), + "indexcache_state": _make_distill_state(0), + } + ) + _tuple_to_dict_helper(pipeline_inputs) + released_inputs = [ + tensor + for tensor in pipeline_inputs + if _get_pipeline_key(tensor) + in { + "indexcache_state 1", + "indexcache_state 2", + "indexcache_state 3", + } + ] + expected_metadata = [ + (list(tensor.shape), tensor.dtype) for tensor in released_inputs + ] + for tensor in released_inputs: + tensor._clear_dataptr() + + hidden_grad = paddle.ones_like(pipeline_inputs[0]) + hidden_grad.stop_gradient = False + input_grads = _normalize_pipeline_input_gradients( + pipeline_inputs, + (hidden_grad, None, None, None), + ) + + self.assertEqual(len(released_inputs), 3) + self.assertEqual(len(input_grads), 4) + for grad, (shape, dtype) in zip(input_grads[1:], expected_metadata): + self.assertEqual(list(grad.shape), shape) + self.assertEqual(grad.dtype, dtype) + self.assertFalse(grad.stop_gradient) + self.assertEqual(float(grad.sum().item()), 0.0) + + meta = SendRecvMeta() + meta.set_send_message(input_grads) + self.assertEqual(len(meta.send_shape_message), 4) + + def test_outer_pipeline_boundary_rejects_missing_hidden_gradient(self): + pipeline_inputs = _dict_to_tuple_helper( + { + "hidden_states": paddle.to_tensor( + [4.0], stop_gradient=False + ), + "indexcache_state": _make_distill_state(0), + } + ) + _tuple_to_dict_helper(pipeline_inputs) + + with self.assertRaisesRegex( + RuntimeError, + "missing a gradient outside IndexCache state", + ): + _normalize_pipeline_input_gradients( + pipeline_inputs, + (None, None, None, None), + ) + + +if __name__ == "__main__": + unittest.main()