From 7c9f0dc355d14a96bef663d89bd7d96624dcf693 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Mon, 21 Sep 2026 22:18:48 +0800 Subject: [PATCH] feat(deepseek_v41): defer activation-recompute decision to the V4.1 backbone Engram's _validate_unsupported_features rejects any activation recomputation, a guard that predates the V4.1 hybrid backbone. That backbone now owns the recompute policy in TransformerConfig._validate_dsv41_config and threads input_ids through _checkpointed_forward for the replayed Engram forward. Add _RecomputeDeferredView, a read-only view that masks only recompute_granularity, and override DeepseekV41EngramConfig._validate_unsupported_features to validate through it. This lets full/selective recompute work for V4.1+Engram while keeping every other Engram unsupported-feature check live and without mutating the shared transformer config. --- src/mcore_bridge/model/modules/engram.py | 28 ++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/src/mcore_bridge/model/modules/engram.py b/src/mcore_bridge/model/modules/engram.py index babd766..8bd6dfe 100644 --- a/src/mcore_bridge/model/modules/engram.py +++ b/src/mcore_bridge/model/modules/engram.py @@ -52,6 +52,29 @@ def __getattr__(self, name): return getattr(self._transformer_config, name) +class _RecomputeDeferredView: + """Read-only view that hides ``recompute_granularity`` from Engram's blanket reject. + + Engram's own ``_validate_unsupported_features`` refuses any activation recomputation, a + guard that predates the V4.1 hybrid backbone. That backbone now owns the recompute policy + in ``TransformerConfig._validate_dsv41_config`` (full-layer replay through the HybridModel + state adapter, or the selective set ``{mhc, layernorm, mla_up_proj, moe_act, moe, + shared_experts}`` with ``mhc`` gated on single-pass Hybrid, and core-attn replay rejected) + and threads ``input_ids`` through ``_checkpointed_forward`` so the replayed Engram forward + still receives its tokens. Masking only the granularity lets the Engram validator keep every + other unsupported-feature check while deferring the recompute decision to that single source + of truth, without mutating the shared transformer config. + """ + + recompute_granularity = None + + def __init__(self, transformer_config): + self._transformer_config = transformer_config + + def __getattr__(self, name): + return getattr(self._transformer_config, name) + + if EngramConfig is not None: class DeepseekV41EngramConfig(EngramConfig): @@ -102,6 +125,11 @@ def _validate_packed_sequences(self, transformer_config, packed_sequences): super()._validate_packed_sequences(transformer_config, packed_sequences) finally: self.variant_spec = variant_spec + + def _validate_unsupported_features(self, transformer_config, use_fsdp): + # Defer the activation-recomputation decision to the V4.1 hybrid backbone (see + # _RecomputeDeferredView); keep every other Engram unsupported-feature check live. + super()._validate_unsupported_features(_RecomputeDeferredView(transformer_config), use_fsdp) else: DeepseekV41EngramConfig = None