[release/0.4][Improvements] Run the DSA warmup phase on dense MHA instead of latent MQA - #1721
Conversation
There was a problem hiding this comment.
| 序号 | 位置 | 优先级 | 状态 |
|---|---|---|---|
| 1 | hybrid_mla_indexer.py:215 |  |
✅ |
| if input_ids is None: | ||
| return None, None | ||
| pad_token_id = getattr(self.config, "pad_token_id", 0) | ||
| assert pad_token_id is not None, ( |
There was a problem hiding this comment.
这里把 pad_token_id=None 直接视为非法配置，但 TransformerConfig.from_config 可以从外部/HF 配置复制该值，仓库的 embedding 和 MoE router 也都约定将 None 回退为 0。warmup 和 sparse 两个新 core 在带 input_ids 的 indexer-loss 训练路径都会走到这个共享函数，因此会在首次计算 loss 时中止；而 python -O 会移除该断言，继续用 None 构造错误的行掩码。请按项目现有约定回退 None 到 0，或在配置入口统一做显式 ValueError 校验，并补上 None 配置的回归测试。
There was a problem hiding this comment.
Good catch — this was wrong, and it is now fixed (fallback to 0, plus a regression test).
Why the assert was there and why it should not have been: the line is a straight copy of
csa_attention.py:2387-2390 (assert pad_token_id is not None, introduced by a87cd90
"[Bug fixes] add loss mask to indexer loss (#1291)"), so it looked like the house style for
this exact spot. It is not — it is the outlier. Every other consumer of the field in this
repository folds None to 0:
gpt_embedding.py:214-216mtp_embedding_layer.py:105-107moe_router.py:605-607,:700-702,:1167-1169,:1226-1228,:1445-1447
So a config whose pad_token_id came back None (e.g. copied out of an external/HF config
by from_config) would be accepted by the embedding and by the router and then abort on the
first indexer loss step of a warmup run. Your second point is just as load-bearing: under
python -O the assert disappears entirely and input_ids != None degenerates to an
all-True mask, i.e. the padding rows silently get counted in both the KL sum and the
denominator — a wrong-number failure instead of a loud one.
What I changed (hybrid_mla_indexer.py, _indexer_loss_mask):
pad_token_id = getattr(self.config, "pad_token_id", 0)
if pad_token_id is None:
pad_token_id = 0with the convention and its seven call sites written into the docstring so the next reader
does not "fix" it back into an assert.
I deliberately did not raise a ValueError at config entry instead. pad_token_id is
consumed by the embedding, the MTP embedding and the router long before any indexer layer
runs, and all three already treat None as 0; making it fatal in TransformerConfig
would change behaviour for models that have nothing to do with the hybrid MLA phases, which
is well outside this PR. Tightening csa_attention.py:2387-2390 the same way is a separate,
CSA-owned change — happy to do it in a follow-up if you want the two indexers to match.
Test (test_hybrid_mla_warmup_doc_mask_loss.py::test_unset_pad_token_id_falls_back_to_zero):
sets config.pad_token_id = None, asserts the mask and the valid-row count are identical
to the pad_token_id=0 module (200 real tokens out of 256), and then drives a full training
step so the fallback is checked where it actually matters — the KL sum and the denominator
the backward divides by — against an independently computed reference (rel. err < 1e-5).
The whole file: 27 passed, 80 subtests.
There was a problem hiding this comment.
已确认当前版本将 None 回退为 0，并补充了覆盖实际 loss 路径的回归测试，该问题已解决。
PaddleFleet Log Analysis
日志分析报告
失败的测试case: 根本原因分析: PR 修复建议:
🔄 每次 Re-run 后自动更新 |
…tead of latent MQA The hybrid-MLA programme has four phases on a ``csa_compress_ratios == -2`` layer: (1) dense MHA pretrain, (2) MHA -> MQA equivalence, (3) DSA warmup that trains the indexer only, (4) sparse MQA + DSA. Phase 3 is *dense* by construction -- the indexer is still learning, so the main attention must not consume its ranking -- yet it ran on ``MQALatentAttention``, i.e. it drove dense semantics through the block-sparse kernel by handing it a full per-document causal index table. That table is O(s^2) per ``-2`` layer and buys nothing: measured peak memory for one core attention (forward + backward) is 8.2x / 8.5x / 9.0x the dense path at s = 1024 / 2048 / 4096, and the ratio grows with s. The absorbed layout also materialises ``q_absorbed`` (``[b, s, h, 576]``). This routes the warmup phase through a dense per-head attention instead, so its memory profile is exactly phase 1's: * ``MHADSAWarmupAttention`` (new) subclasses ``DotProductAttention`` and delegates the whole attention half to ``super().forward``, adding only the indexer projections and the KL loss. Phase 2 is therefore "phase 1 plus an indexer loss" by construction rather than by assertion, and the q/k it produces are bit-identical to phase 1's (verified with ``apply_rope_fusion`` both off and on). * ``hybrid_mla_indexer.py`` (new) holds the one dispatch predicate, ``latent_mqa_enabled(config)``, plus the indexer plumbing both backends share. ``gpt_layer_specs`` and ``MultiLatentAttention`` are the only callers, so no code path can pick a backend the model would not build. ``dsa_indexer_use_sparse_loss`` now selects the backend class, not just the KL width, and must be passed as a construction kwarg. * ``MQALatentAttention`` keeps only the sparse phases; the full-causal index table construction it needed for the warmup is gone (-434 lines net there). Parameter names, ``state_dict`` keys and saved HF keys are unchanged in every phase (16 keys; phase 1 -> warmup adds the 5 ``indexer.*`` keys and nothing else), so checkpoints stay loadable across the switch. Two fixes fell out of running the new backend: * ``TileLangCSAIndexerLossAutoScaler`` returned its first argument unchanged, and Paddle records a PyLayer returning one of its inputs as an inplace write on it. The dense attention backward saves its own output, so the version bump made the *attention* backward raise ``PermissionDenied: Tensor ... modified by an inplace operation``. The caller now hands the scaler a fresh tensor (``clone`` is a gradient identity). This was invisible before because the scaler already clones when the backbone is frozen. * Three context-parallel cases asserted a bit-identical ``dq`` against phase 1. Dense flashmask accumulates ``dq`` atomically over column blocks, so two runs of the *same* module differ by 3.05e-05 at s = 512 while forward, ``dk`` and ``dv`` are exact. The bound is now self-calibrated from a second reference run measured in the same test run. Rebased on PaddlePaddle#1679, which adds ``mqa_split_kv_b_proj``. That switch replaces ``kv_b_proj`` with standalone ``k_b_proj`` / ``v_b_proj``, so it only means anything where absorption happens. The warmup phase is now one of the dense phases and keeps ``kv_b_proj``, so its validation was moved onto ``latent_mqa_enabled`` too: the combination is rejected at config time instead of silently changing the parameter set at the warmup -> sparse switch. Tested on SM100 after the rebase: 431 single-card tests over 14 files (config pipeline 24, doc equivalence 35, dsv4 hybrid 87, grad health 22, HF roundtrip 12, latent MQA 63, Muon 67, warmup RoPE/recompute/MTP 14, MLA RoPE CP + VHA 56, warmup doc-mask loss + train_indexer_only 51) and 30 two-card CP tests (test_mqa_dsa_cp 8, test_mqa_dsa_warmup_cp 8, test_indexer_topk_col_mask_cp 4, test_mla_cp_contiguous_allgather 10). The one failure, ``test_documented_bug_gate_proj_is_saved_untransposed``, predates this branch: it pins a defect that has since been fixed outside this repository, and it fails identically on release/0.4.
…r loss mask Review of PaddlePaddle#1721 flagged the `assert pad_token_id is not None` copied into `HybridMLAIndexerMixin._indexer_loss_mask` from `csa_attention.py:2387-2390` (upstream PaddlePaddle#1291): `TransformerConfig.from_config` can copy a `None` straight out of an external/HF config, so a run the embedding and the MoE router both accept would abort at its first indexer loss -- and under `python -O`, where asserts are stripped, the ids would have been compared against `None` and no row masked at all. Fold `None` to `0`, the convention every other consumer of the field already follows (`gpt_embedding.py:214-216`, `mtp_embedding_layer.py:105-107`, `moe_router.py:605-607` and four more sites), and pin it with a regression test that asserts the `None` config produces the same mask, the same denominator and the same logged loss as `pad_token_id=0`. The line references in the two warmup tests move with the edit.
…the phase-2 drift sentinel `TestConfigDeltas` pins, key by key, how each experiment YAML may differ from the phase-1 baseline, and it fails both ways: an unlisted difference fails, and an allowlisted key that stops differing fails too (so the allowlist cannot rot into a blanket exemption). Now that the warmup phase runs the baseline's dense attention -- its `-2` layers build `MHADSAWarmupAttention`, a `DotProductAttention` subclass -- its memory/compute picture is phase 1's, and the phase-2 YAML was resynced to the baseline's `recompute_granularity: selective` + module list and `gradient_accumulation_steps: 2`. Those five keys therefore no longer differ for that one config and the sentinel's second assertion fires. Remove them from its allowlist; the three latent-MQA variants keep them, because their picture really does differ. Only reachable where the parent config repository is checked out above PaddleFleet (the YAMLs live there), so CI skips it.
eb44594 to
ea48c10
Compare
Codecov Report❌ Patch coverage is ❌ Your patch status has failed because the patch coverage (70.22%) 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.4 #1721 +/- ##
==============================================
Coverage ? 70.22%
==============================================
Files ? 6
Lines ? 178
Branches ? 27
==============================================
Hits ? 125
Misses ? 46
Partials ? 7
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
4 participants
PR Category
Operator Mechanism
PR Types
Improvements
Description
The hybrid-MLA programme has four phases on a
csa_compress_ratios == -2layer:(1) dense MHA pretrain, (2) MHA -> MQA equivalence, (3) DSA warmup that trains
the indexer only, (4) sparse MQA + DSA. Phase 3 is dense by construction --
the indexer is still learning, so the main attention must not consume its
ranking -- yet it ran on
MQALatentAttention, i.e. it drove dense semanticsthrough the block-sparse kernel by handing it a full per-document causal index
table. That table is
O(s^2)per-2layer and buys nothing.Measured peak memory for one core attention (forward + backward, bf16, same
fixture geometry, only the attention half differing):
The ratio grows with
s, and it is a conservative lower bound: the control armcarries no indexer while the dense arm does.
This routes the warmup phase through dense per-head attention, so its memory
profile is exactly phase 1's:
MHADSAWarmupAttention(new) subclassesDotProductAttentionanddelegates the whole attention half to
super().forward, adding only theindexer projections and the KL loss. Phase 2 is therefore "phase 1 plus an
indexer loss" by construction rather than by assertion, and the q/k it
produces are bit-identical to phase 1's (verified with
apply_rope_fusionboth off and on:
q/kmaxabs0.000e+00).hybrid_mla_indexer.py(new) holds the one dispatch predicate,latent_mqa_enabled(config), plus the indexer plumbing both backends share.gpt_layer_specsandMultiLatentAttentionare the only callers, so no codepath can pick a backend the model would not build.
dsa_indexer_use_sparse_lossnow selects the backend class, not just the KLwidth, and must be passed as a construction kwarg.
MQALatentAttentionkeeps only the sparse phases; the full-causal indextable construction it needed for the warmup is gone (-434 lines net there).
Parameter names,
state_dictkeys and saved HF keys are unchanged in everyphase (16 keys; phase 1 -> warmup adds the 5
indexer.*keys and nothing else),so checkpoints stay loadable across the switch.
Two fixes fell out of running the new backend:
TileLangCSAIndexerLossAutoScalerreturned its first argument unchanged, andPaddle records a PyLayer returning one of its inputs as an inplace write on
it. The dense attention backward saves its own output, so the version bump
made the attention backward raise
PermissionDenied: Tensor ... modified by an inplace operation. The callernow hands the scaler a fresh tensor (
cloneis a gradient identity). This wasinvisible before because the scaler already clones when the backbone is
frozen.
dqagainst phase 1.Dense flashmask accumulates
dqatomically over column blocks, so two runs ofthe same module differ by
3.05e-05ats = 512while forward,dkanddvare exact. The bound is now self-calibrated from a second reference runmeasured in the same test run.
Rebased on #1679, which adds
mqa_split_kv_b_proj. That switch replaceskv_b_projwith standalonek_b_proj/v_b_proj, so it only means anythingwhere absorption happens. The warmup phase is now one of the dense phases and
keeps
kv_b_proj, so its validation was moved ontolatent_mqa_enabledtoo:the combination is rejected at config time instead of silently changing the
parameter set at the warmup -> sparse switch.
Tested
On SM100, after the rebase: 431 single-card tests over 14 files (config
pipeline 24, doc equivalence 35, dsv4 hybrid 87, grad health 22, HF roundtrip
12, latent MQA 63, Muon 67, warmup RoPE/recompute/MTP 14, MLA RoPE CP + VHA 56,
warmup doc-mask loss +
train_indexer_only51) and 30 two-card CP tests(
test_mqa_dsa_cp8,test_mqa_dsa_warmup_cp8,test_indexer_topk_col_mask_cp4,test_mla_cp_contiguous_allgather10).The one failure,
test_documented_bug_gate_proj_is_saved_untransposed, predatesthis branch: it pins a defect that has since been fixed outside this
repository, and it fails identically on
release/0.4.