From 4ed70de25e7f369eb3fe68ec64616ee80e63205c Mon Sep 17 00:00:00 2001 From: huangjiyi Date: Thu, 30 Jul 2026 15:07:34 +0800 Subject: [PATCH] [KDA] Eagerly load Paddle training imports --- fla/modules/conv/cp/ops.py | 18 +++--- fla/modules/conv/short_conv.py | 10 ++- tests/paddle/test_kimi_delta_attention.py | 75 +++++++++++++++++++---- 3 files changed, 79 insertions(+), 24 deletions(-) diff --git a/fla/modules/conv/cp/ops.py b/fla/modules/conv/cp/ops.py index c18534ff13..ba0445ff4b 100644 --- a/fla/modules/conv/cp/ops.py +++ b/fla/modules/conv/cp/ops.py @@ -8,8 +8,12 @@ import torch import torch.distributed as dist +# initialize FLA ops before the convolution backend to avoid the modules/ops import cycle. +# isort: off from fla.ops.cp import FLACPContext, conv_cp_send_recv_bwd, conv_cp_send_recv_fwd from fla.ops.utils import prepare_chunk_indices +from fla.modules.conv.triton.ops import causal_conv1d_bwd, causal_conv1d_fwd +# isort: on class CausalConv1dFunctionCP(torch.autograd.Function): @@ -129,9 +133,6 @@ def forward( chunk_size: int | None, backend: str = 'triton', ): - # Import here to avoid circular dependency - from fla.modules.conv.triton.ops import causal_conv1d_fwd - if cp_context is None: raise ValueError("cp_context must be provided for CausalConv1dFunctionCP") cu_seqlens = cp_context.cu_seqlens @@ -150,6 +151,8 @@ def forward( ) ctx.save_for_backward(x, weight, bias, initial_state) + ctx.has_bias = bias is not None + ctx.has_chunk_indices = chunk_indices is not None ctx.activation = activation ctx.cu_seqlens = cu_seqlens ctx.cu_seqlens_cpu = cu_seqlens_cpu @@ -179,9 +182,6 @@ def forward( @staticmethod def backward(ctx, dy: torch.Tensor): - # Import here to avoid circular dependency - from fla.modules.conv.triton.ops import causal_conv1d_bwd - x, weight, bias, initial_state = ctx.saved_tensors group = ctx.group W = ctx.W @@ -212,7 +212,11 @@ def backward(ctx, dy: torch.Tensor): pre_num_conv_tokens=ctx.pre_num_conv_tokens, ) - return dx, dw, db, None, None, None, None, None + return ( + (dx, dw) + + ((db,) if ctx.has_bias else ()) + + ((None,) if ctx.has_chunk_indices else ()) + ) def causal_conv1d_cp( diff --git a/fla/modules/conv/short_conv.py b/fla/modules/conv/short_conv.py index 2509e29189..254159c37e 100644 --- a/fla/modules/conv/short_conv.py +++ b/fla/modules/conv/short_conv.py @@ -9,12 +9,16 @@ from __future__ import annotations +import os import warnings import torch import torch.nn as nn from einops import rearrange +from fla.modules.conv.causal_conv1d import causal_conv1d +from fla.modules.conv.triton.ops import causal_conv1d_update + try: from causal_conv1d import causal_conv1d_fn as causal_conv1d_fn_cuda from causal_conv1d import causal_conv1d_update as causal_conv1d_update_cuda @@ -85,7 +89,6 @@ def __init__( "The `use_fast_conv1d` parameter is deprecated and will be ignored. " "Please use the `backend` parameter instead.", ) - import os self.backend = os.environ.get('FLA_CONV_BACKEND', backend) if backend not in ['cuda', 'triton']: raise ValueError(f"Invalid backend: {backend}, must be one of ['cuda', 'triton']") @@ -151,9 +154,6 @@ def forward( Returns: Tensor of shape `[B, T, D]`. """ - # Import here to avoid circular dependency - from fla.modules.conv.causal_conv1d import causal_conv1d - B, T, *_ = x.shape N = B if cu_seqlens is None else len(cu_seqlens) - 1 if mask is not None: @@ -208,8 +208,6 @@ def step( output_final_state: bool = False, cu_seqlens: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - from fla.modules.conv.triton.ops import causal_conv1d_update - B, _, D, W = *x.shape, self.kernel_size[0] N = B if cu_seqlens is None else len(cu_seqlens) - 1 # Always initialise cache when None so the Triton kernel never diff --git a/tests/paddle/test_kimi_delta_attention.py b/tests/paddle/test_kimi_delta_attention.py index 0017a02a72..ad6608dce5 100644 --- a/tests/paddle/test_kimi_delta_attention.py +++ b/tests/paddle/test_kimi_delta_attention.py @@ -5,8 +5,10 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors +import builtins import subprocess import sys +from contextlib import contextmanager from types import SimpleNamespace import paddle @@ -19,11 +21,35 @@ paddle.disable_compat() from fla.modules import FusedRMSNormGated, ShortConvolution # noqa: E402 +from fla.ops.cp import FLACPContext # noqa: E402 from fla.ops.kda import chunk_kda # noqa: E402 from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask # noqa: E402 from fla.utils import tensor_cache # noqa: E402 +@contextmanager +def _assert_no_fla_runtime_imports(): + loaded_modules = {name for name in sys.modules if name == "fla" or name.startswith("fla.")} + runtime_imports = [] + original_import = builtins.__import__ + + def monitored_import(name, globals_=None, locals_=None, fromlist=(), level=0): + caller = globals_.get("__name__", "") if globals_ else "" + if caller == "fla" or caller.startswith("fla."): + runtime_imports.append((caller, name, tuple(fromlist or ()), level)) + return original_import(name, globals_, locals_, fromlist, level) + + builtins.__import__ = monitored_import + try: + yield + finally: + builtins.__import__ = original_import + + assert runtime_imports == [] + current_modules = {name for name in sys.modules if name == "fla" or name.startswith("fla.")} + assert current_modules == loaded_modules + + @tensor_cache def _get_unpad_data(attention_mask: paddle.Tensor) -> tuple[paddle.Tensor, paddle.Tensor]: indices = paddle.nonzero(attention_mask.flatten()).flatten() @@ -249,6 +275,32 @@ def test_fused_rms_norm_gated_forward_backward(): paddle.testing.assert_close(layer.weight.grad, weight_ref.grad, rtol=1e-5, atol=1e-5) +def test_short_convolution_cp_forward_backward_uses_eager_imports(): + paddle.seed(42) + conv = ShortConvolution(64, 4, activation="silu") + x = paddle.randn([1, 32, 64], dtype="float32") + x.stop_gradient = False + cu_seqlens = paddle.to_tensor([0, 32], dtype="int32") + cp_context = FLACPContext( + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=paddle.to_tensor([0, 32], dtype="int32", place=paddle.CPUPlace()), + is_first_rank=True, + conv1d_kernel_size=4, + pre_num_conv_tokens=0, + ) + + with _assert_no_fla_runtime_imports(): + output, final_state = conv(x, cp_context=cp_context) + output.square().mean().backward() + + assert final_state is None + assert paddle.isfinite(output).all().item() + assert x.grad is not None + assert paddle.isfinite(x.grad).all().item() + assert conv.weight.grad is not None + assert paddle.isfinite(conv.weight.grad).all().item() + + @pytest.mark.parametrize("use_padding_mask", [False, True]) def test_kimi_delta_attention_training_forward_backward(use_padding_mask: bool): paddle.seed(42) @@ -262,17 +314,18 @@ def test_kimi_delta_attention_training_forward_backward(use_padding_mask: bool): "gate_lower_bound": -5.0, }, ) - layer = _KimiDeltaAttentionTrainingHarness(config) - hidden_states = paddle.randn([2 if use_padding_mask else 1, 64, 128], dtype="float32") - hidden_states.stop_gradient = False - attention_mask = None - if use_padding_mask: - attention_mask = paddle.to_tensor([[1] * 53 + [0] * 11, [1] * 64], dtype="bool") - - with paddle.amp.auto_cast(enable=True, dtype="bfloat16"): - output = layer(hidden_states, attention_mask) - loss = output.astype("float32").square().mean() - loss.backward() + with _assert_no_fla_runtime_imports(): + layer = _KimiDeltaAttentionTrainingHarness(config) + hidden_states = paddle.randn([2 if use_padding_mask else 1, 64, 128], dtype="float32") + hidden_states.stop_gradient = False + attention_mask = None + if use_padding_mask: + attention_mask = paddle.to_tensor([[1] * 53 + [0] * 11, [1] * 64], dtype="bool") + + with paddle.amp.auto_cast(enable=True, dtype="bfloat16"): + output = layer(hidden_states, attention_mask) + loss = output.astype("float32").square().mean() + loss.backward() assert output.shape == hidden_states.shape assert paddle.isfinite(output).all().item()