[release/0.3] fix fleet fa version dispatch

[risemeup1111]

PR Category

Distributed Strategy

PR Types

Bug fixes

Description

fix fa version dispatch in fleet

是否引起精度变化

否

Cherry-pick of #1396 to release/0.3.

Merged in dev: #1396

[PaddlePaddle-bot]

🤖 Paddle-CI-Agent | pr_review | 2026-07-04 17:54:37

📋 Review 摘要

PR 概述：修正 Fleet FlashAttention/FlashMask 在 fa_version 与 deterministic 场景下的版本分发。
变更范围：context_parallel_utils.py CP FlashMask 前后向版本传递、refined_recompute/flash_attn.py recompute FlashAttention 版本选择。
影响面 Tag：Context Parallel Recompute

问题

级别	文件	概述
🔴 Bug	src/paddlefleet/refined_recompute/flash_attn.py:67	v4 不可用回退到 v3 后未重新应用 deterministic fallback
🔴 Bug	src/paddlefleet/context_parallel_utils.py:733	CP 前向在 v4 后端不可用时可能返回与实际 kernel 不一致的 fa_version

📝 PR 规范检查

标题不符合 PaddleFleet Cherry-Pick 标题格式（当前使用 [release/0.3]，不是 checklist §D1 中的规范 Tag）。

标题建议（可直接复制）：

• [Cherry-Pick][Bug fixes] fix fleet fa version dispatch (#1396)

PR 描述建议（点击展开，可直接复制）

### PR Category
<!-- One of [ User Experience | Execute Infrastructure | Operator Mechanism | Custom Device | Performance Optimization | Distributed Strategy | Parameter Server | Communication Library  | Environment Adaptation ] -->
Distributed Strategy
### PR Types
<!-- One of [ New features | Bug fixes | Improvements | Performance | BC Breaking | Deprecations | Docs | Devs | Not User Facing | Security | Others ] -->
Bug fixes
### Description
<!-- Describe what you’ve done -->
fix fa version dispatch in fleet

Cherry-pick of #1396 to `release/0.3`.

Merged in dev: #1396  <!-- For pass CI -->

总体评价

当前修改方向合理，但两个路径都把 deterministic fallback 绑定在原始 flag 上，漏掉 v4 不可用后的有效版本回退场景。建议先统一计算 effective fa_version，再按 effective version 做 deterministic 降级并传递给后向。

[PaddlePaddle-bot]

🔴 Bug 这里用原始 flag 是否等于 3 来决定 deterministic fallback，会漏掉 FLAGS_flash_attn_version=4 但 _flash_mask_available=False 的回退路径。

当前函数后面会把不可用的 v4 回退成 v3；但这个回退发生在 deterministic 检查之后。确定性模式且 hdim > 128（或旧签名无 block_mask）时，原逻辑会返回 v2，现在会返回 v3，违背上面的 deterministic 不支持约束，并可能在 recompute 的前/后向中调用不支持的 v3 kernel。

建议修复方式：
先把用户请求的版本归一化为有效版本，再对有效的 v3 执行 deterministic fallback，例如：

if fa_version == 4 and not _flash_mask_available:
    fa_version = 3
if fa_version == 3 and deterministic_needs_v2:
    return 2
return fa_version

[PaddlePaddle-bot]

🔴 Bug CP 路径同样只检查原始 fa_version == 3，导致请求 v4 但 _flash_mask_available=False 时前向实际走普通 flashmask_attention，却仍把 fa_version=4 返回给后向。

cp_flashmask_allgatherkv_balance_backward() 只有在 fa_version == 4 and _flash_mask_available 时才走 v4 分支，否则会落到 ValueError("FlashAttention version 4 is not supported.")。确定性配置下旧逻辑会降到 v2，新逻辑跳过降级后会让 CP 训练在 backward 阶段失败。

建议修复方式：
像 _get_fa_version() 一样集中计算 effective version；当 v4 后端不可用时把 fa_version 改成普通 FlashMask 实际使用的 3/2，并把这个值返回给 backward。

[codecov-commenter]

Codecov Report

❌ Patch coverage is 54.54545% with 5 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (release/0.3@5de515c). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
src/paddlefleet/refined_recompute/flash_attn.py	16.66%	3 Missing and 2 partials ⚠️

❌ Your patch status has failed because the patch coverage (54.54%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@              Coverage Diff               @@
##             release/0.3    #1398   +/-   ##
==============================================
  Coverage               ?   54.54%           
==============================================
  Files                  ?        2           
  Lines                  ?       11           
  Branches               ?        8           
==============================================
  Hits                   ?        6           
  Misses                 ?        3           
  Partials               ?        2           

Flag	Coverage Δ
coverage_combine	54.54% <54.54%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/paddlefleet/context_parallel_utils.py	100.00% <100.00%> (ø)
src/paddlefleet/refined_recompute/flash_attn.py	16.66% <16.66%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Paddle-CI-Bot]

PaddleFleet Log Analysis

Run #28783919413 · Attempt 1

日志分析报告

流水线名称	问题标签	修复建议	日志片段
Unit test (single card)	单测超时（60 min）	test_dsv4_hybrid_attention.py 在运行 8 个用例后发生 hang，需要排查该用例是否存在无限循环或死锁	报错代码
Coverage Upload And Check	增量覆盖率不达标（exit 9）	src/paddlefleet/refined_recompute/flash_attn.py 第 69、75、79 行未被覆盖，需补充对应单测	报错代码

失败的测试case:

tests/single_card_tests/transformer/test_dsv4_hybrid_attention.py  （超时 hang，60min 未完成）

根本原因分析:
PR cherry-pick/1396 引入了 DSv4 hybrid attention 相关改动，test_dsv4_hybrid_attention.py 在执行第 8 个用例时陷入 hang（无输出、无完成），导致 60min 超时；同时 refined_recompute/flash_attn.py 新增代码（line 69/75/79）缺乏单测覆盖，增量覆盖率跌至 72%，低于 90% 阈值。

修复建议:

1. 超时 hang：本地复现 test_dsv4_hybrid_attention.py 中第 8 个之后的用例（run 时加 -v -s 观察卡在哪个 case）；排查 DSv4 hybrid attention forward/backward 中是否有等待通信/锁未释放的路径，重点检查本 PR 改动的算子调用链。

2. 增量覆盖率：在 tests/single_card_tests/ 下为 refined_recompute/flash_attn.py 的 line 69（条件分支）、75（另一分支）、79（边界处理）补充对应 pytest 用例，确保 diff coverage ≥ 90%。

---

🔍 准确性记录：请点击评论底部 😊 图标，选择 👍（准确）或 👎（有误），将自动记录到 CI 监控系统

_{🔄 每次 Re-run 后自动更新}