Skip to content

fix fleet fa version dispatch #1396

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tianlef merged 1 commit into from
Jul 4, 2026
Merged

fix fleet fa version dispatch #1396

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions src/paddlefleet/context_parallel_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -730,11 +730,12 @@ def cp_flashmask_allgatherkv_balance_forward(
deterministic = paddle.get_flags(["FLAGS_cudnn_deterministic"])[
"FLAGS_cudnn_deterministic"
]
if "block_mask" in inspect.signature(flashmask_attention).parameters:
if deterministic and query.shape[-1] > 128:
if fa_version == 3:

@PaddlePaddle-bot PaddlePaddle-bot Jul 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Bug 这里同样只处理原始 fa_version == 3。当 FLAGS_flash_attn_version=4 但 _flash_mask_available=False 时，forward 实际不会进入 v4 _flash_attn_fwd，却仍把 fa_version=4 返回给 FlashMaskContextParallel.backward。

随后 cp_flashmask_allgatherkv_balance_backward 的 v4 分支又要求 _flash_mask_available 为真；在 deterministic 训练场景下，这次改动会跳过原本的 v2 fallback，并在 backward 分发到 unsupported version。

建议修复方式：在 forward 中先计算 effective_fa_version。如果 v4 backend 不可用，应先降为 3，再按 deterministic 规则必要时降为 2，并把这个实际使用的版本返回给 backward。

if "block_mask" in inspect.signature(flashmask_attention).parameters:
if deterministic and query.shape[-1] > 128:
fa_version = 2
elif deterministic:
fa_version = 2
elif deterministic:
fa_version = 2

if fa_version == 4 and _flash_mask_available:
Comment on lines +733 to 740

@risemeup1111 risemeup1111 Jul 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级：P1
这里用原始 flag 直接判断 fa_version == 3，会漏掉 FLAGS_flash_attn_version=4 但 _flash_mask_available=False 的路径。该路径实际会走下面的 flashmask_attention(...) fallback，但函数仍返回 fa_version=4 给 backward；cp_flashmask_allgatherkv_balance_backward 只有 fa_version == 4 and _flash_mask_available 分支，未安装 cute backend 时会落到 ValueError。请先把不可用的 FA4 规范化到有效 fallback 版本，再做 deterministic 降级并返回这个有效版本。

Suggested change
if fa_version == 3:
if "block_mask" in inspect.signature(flashmask_attention).parameters:
if deterministic and query.shape[-1] > 128:
fa_version = 2
elif deterministic:
fa_version = 2
elif deterministic:
fa_version = 2
if fa_version == 4 and _flash_mask_available:
if fa_version == 4 and not _flash_mask_available:
fa_version = 3
if fa_version == 3:
if "block_mask" in inspect.signature(flashmask_attention).parameters:
if deterministic and query.shape[-1] > 128:
fa_version = 2
elif deterministic:
fa_version = 2
if fa_version == 4:

@umiswing umiswing Jul 4, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

此处应该直接报错，提醒用户H卡上不要指定使用FA4

@risemeup1111 risemeup1111 Jul 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

可以，直接报错也能解决这里的问题，而且语义更明确：用户显式指定 FLAGS_flash_attn_version=4 时，如果当前 H 卡/环境没有 cute FA4 backend，就不应该静默走 flashmask_attention(...) 并把 fa_version=4 传给 backward。关键是要在进入 fallback 调用前拦住，例如：

if fa_version == 4 and not _flash_mask_available:
    raise NotImplementedError(
        "FA4 FlashMask requires cute backend; do not set FLAGS_flash_attn_version=4 on H-card/current environment."
    )

如果采用这个方案，后面 if fa_version == 4 and _flash_mask_available: 也可以相应简化成 if fa_version == 4:，并确保函数不会再返回一个实际未使用的 fa_version=4。

output, log_sum_exp = _flash_attn_fwd(
Expand Down
25 changes: 13 additions & 12 deletions src/paddlefleet/refined_recompute/flash_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,21 +61,22 @@ def _get_fa_version(hdim):
if "xpu" in paddle.get_device():
return 2
# Xiangrui: For deterministic, NOT support for hdim > 128 currently.
if "block_mask" in inspect.signature(flashmask_attention).parameters:
if (
paddle.get_flags(["FLAGS_cudnn_deterministic"])[
"FLAGS_cudnn_deterministic"
]
and hdim > 128
):
return 2
elif paddle.get_flags(["FLAGS_cudnn_deterministic"])[
"FLAGS_cudnn_deterministic"
]:
return 2
fa_version = paddle.base.framework.get_flags(["FLAGS_flash_attn_version"])[
"FLAGS_flash_attn_version"
]
if fa_version == 3:

@PaddlePaddle-bot PaddlePaddle-bot Jul 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Bug 这里只在原始 FLAGS_flash_attn_version == 3 时执行 deterministic 降级，但下面 fa_version == 4 and not _flash_mask_available 会把有效版本回退成 3。

在未安装/不支持 cute FlashMask 且用户设置 FLAGS_flash_attn_version=4 时，deterministic + hdim > 128（或旧签名无 block_mask）会直接返回 3，后续 refined recompute 会走 v3 kernel，绕过此前的 v2 fallback。

建议修复方式：先把 flag 解析成实际可用版本，再对有效 v3 套 deterministic 降级，例如：

if fa_version == 4 and not _flash_mask_available:
    logger.warning(...)
    fa_version = 3
if fa_version == 3 and deterministic_fallback_needed:
    return 2
return fa_version

if "block_mask" in inspect.signature(flashmask_attention).parameters:
if (
paddle.get_flags(["FLAGS_cudnn_deterministic"])[
"FLAGS_cudnn_deterministic"
]
and hdim > 128
):
return 2
elif paddle.get_flags(["FLAGS_cudnn_deterministic"])[
"FLAGS_cudnn_deterministic"
]:
return 2
# Fall back to version 3 if flash_mask is not available
if fa_version == 4 and not _flash_mask_available:
Comment on lines +67 to 81

@risemeup1111 risemeup1111 Jul 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级：P1
fa_version == 4 and not _flash_mask_available 时当前代码会在这个 block 之后直接 return 3，因此没有再应用 deterministic fallback。结果是在用户打开 deterministic、flashmask_attention 支持 block_mask 且 hdim > 128 时，原本应该避开不支持的 FA3 大 head_dim deterministic 路径，却会返回 3 并进入 _C_ops.flash_attn_v3/flashmask_attention_v2。请先把不可用的 FA4 规范化成 3，再复用同一段 deterministic 降级逻辑。

Suggested change
if fa_version == 3:
if "block_mask" in inspect.signature(flashmask_attention).parameters:
if (
paddle.get_flags(["FLAGS_cudnn_deterministic"])[
"FLAGS_cudnn_deterministic"
]
and hdim > 128
):
return 2
elif paddle.get_flags(["FLAGS_cudnn_deterministic"])[
"FLAGS_cudnn_deterministic"
]:
return 2
# Fall back to version 3 if flash_mask is not available
if fa_version == 4 and not _flash_mask_available:
if fa_version == 4 and not _flash_mask_available:
logger.warning(
"FlashMask (fa_version=4) is not available, falling back to fa_version=3"
)
fa_version = 3
if fa_version == 3:
if "block_mask" in inspect.signature(flashmask_attention).parameters:
if (
paddle.get_flags(["FLAGS_cudnn_deterministic"])[
"FLAGS_cudnn_deterministic"
]
and hdim > 128
):
return 2
elif paddle.get_flags(["FLAGS_cudnn_deterministic"])[
"FLAGS_cudnn_deterministic"
]:
return 2

@umiswing umiswing Jul 4, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

此处应该直接报错，提醒用户在H卡上不要指定使用FA4

@risemeup1111 risemeup1111 Jul 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

可以，直接报错也可以作为修复方案。这里的问题不在于一定要 fallback 到 3，而是当前 fa_version=4 && !_flash_mask_available 会继续 return 3，从而绕过 deterministic 降级并进入不期望的 FA3 路径。

如果预期是提醒用户 H 卡上不要指定 FA4，建议把这个分支改成明确异常，而不是 warning 后 return 3：

if fa_version == 4 and not _flash_mask_available:
    raise NotImplementedError(
        "FA4 FlashMask requires cute backend; do not set FLAGS_flash_attn_version=4 on H-card/current environment."
    )

这样也能消除我前面指出的 fallback 后继续执行 FA3/后向分发风险。

logger.warning(
Expand Down
Loading