fix(pi052): accept use_flex_attention config field for checkpoint compat

Newer PI052 training runs serialize use_flex_attention into config.json.
This branch's attention path is SDPA/eager (mathematically equivalent), so
the field is accepted as an inert no-op (mirrors the existing use_hf_kernels
compat field) — otherwise loading those checkpoints raises DecodingError.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
pepijn223
2026-07-13 16:46:10 +02:00
parent 33f0414733
commit 9f0e4dfb53
@@ -246,6 +246,15 @@ class PI052Config(PI05Config):
``DecodingError: The fields use_hf_kernels are not valid for
PI052Config`` (job 22164492). Remove in a future major bump."""
use_flex_attention: bool = False
"""Accepted for checkpoint-config compatibility only — no-op in this branch.
Newer training runs serialize ``use_flex_attention`` into ``config.json`` to
select the FlexAttention kernel at train time; this branch's attention path
is SDPA/eager, which is mathematically equivalent (same softmax attention,
different kernel), so inference/eval results are unchanged. Retained so those
checkpoints load instead of raising ``DecodingError: The fields
use_flex_attention are not valid for PI052Config``."""
# Optimizer foreach/fused. pi052 carries these locally because the shared
# PI05Config (kept identical to upstream main) does not define them; the
# checkpoints we train serialize both keys into config.json, so they must