Commit Graph

2 Commits

Author SHA1 Message Date
junran
4011f5e3d5 Fix torch.compile graph breaks with Engram module
Refactor Engram data structures for torch.compile compatibility:
- Replace nn.ModuleDict (string-keyed) with nested nn.ModuleList (integer-indexed)
- Precompute hash powers in __init__ and re-materialize in init_weights to
  eliminate torch.tensor() construction during forward
- Remove hash_seeds and table_sizes dicts (info embedded in _powers / embed_tables)
- Add _patch_removed_state_keys for old engram checkpoint backward compat

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-28 08:18:50 +08:00
junran
dcd9b0668b Add Engram conditional memory module
Integrates DeepSeek's Engram (N-gram hash lookup + context-aware gating
+ depthwise causal conv) as an optional module behind --engram CLI flag.
Placed at two layers per paper ablation findings (layer 1 and n_layer//2-1).
Coexists with existing Value Embeddings; disabled by default.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-19 12:04:02 +08:00