fix: [gemma4] fix VRAM leak in hybrid FA2+SDPA (hybrid attentiuon) path under activation check…#3611

Merged

winglian merged 6 commits into

axolotl-ai-cloud:mainfrom

thad0ctor:fix/gemma4-hybrid-vram-leak

Apr 21, 2026

Commits on Apr 21, 2026

[gemma4] fix VRAM leak in hybrid FA2+SDPA path under activation checkpointing

thad0ctor
authored and
winglian
committed
Coderabbit comment: gemma4: clear TLS unconditionally in decoder-layer patched __call__

thad0ctor
authored and
winglian
committed
fix: swap threading.local() for module-level store so autograd worker threads see shared_kv_states during backward recompute

thad0ctor
authored and
winglian
committed
scope gemma4 shared_kv_states side channel to checkpointed training

thad0ctor
authored and
winglian
committed
add gemma4 cross-thread visibility test for shared_kv_states store

thad0ctor
authored and
winglian
committed
fix logger
winglian
committed