Skip to content

Commit fec9978

Browse files
authored
fix: activation memory config crashes get_default() with OpenAI (MemTensor#1051)
## Summary Two bugs in `default_config.py` that crash `get_default()` when `enable_activation_memory=True`: **Bug 1: `act_mem` injected into MOSConfig (extra=forbid)** `get_default_config()` adds `act_mem` dict to the config passed to `MOSConfig(**config_dict)`, but `MOSConfig` has no `act_mem` field and inherits `extra="forbid"` from `BaseConfig`: ``` pydantic_core.ValidationError: 1 validation error for MOSConfig act_mem Extra inputs are not permitted [type=extra_forbidden] ``` **Bug 2: KV cache with `openai` backend** `get_default_cube_config()` hardcodes `extractor_llm` backend to `"openai"` for KV cache memory, but `KVCacheMemoryConfig` validator requires `huggingface`/`huggingface_singleton`/`vllm` (KV cache extracts internal attention tensors via `build_kv_cache`, which needs local model access): ``` ConfigurationError: KVCacheMemoryConfig requires extractor_llm backend to be 'huggingface' or 'huggingface_singleton', got 'openai' ``` ## Reproduction ```python from memos.mem_os.utils.default_config import get_default # This crashes with ValidationError config, cube = get_default( openai_api_key="sk-...", enable_activation_memory=True, ) ``` ## Fix - Remove `act_mem` from `get_default_config()` — the `enable_activation_memory` bool flag is sufficient for `MOSConfig`; `act_mem` config belongs in MemCube config only. - In `get_default_cube_config()`, require explicit `activation_memory_backend` kwarg (e.g. `"huggingface"`) instead of hardcoding `"openai"`. Skip `act_mem` gracefully when no compatible local backend is provided. ## Test plan - [ ] `get_default(openai_api_key="sk-...", enable_activation_memory=True)` no longer crashes - [ ] `get_default(openai_api_key="sk-...", enable_activation_memory=True, activation_memory_backend="huggingface", activation_memory_llm_config={...})` creates valid act_mem config - [ ] `get_default(openai_api_key="sk-...")` continues to work as before
2 parents d3dc6b8 + e260982 commit fec9978

File tree

1 file changed

+30
-26
lines changed

1 file changed

+30
-26
lines changed

src/memos/mem_os/utils/default_config.py

Lines changed: 30 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,15 @@
33
Provides simplified configuration generation for users.
44
"""
55

6+
import logging
67
from typing import Literal
78

89
from memos.configs.mem_cube import GeneralMemCubeConfig
910
from memos.configs.mem_os import MOSConfig
1011
from memos.mem_cube.general import GeneralMemCube
1112

13+
logger = logging.getLogger(__name__)
14+
1215

1316
def get_default_config(
1417
openai_api_key: str,
@@ -116,20 +119,9 @@ def get_default_config(
116119
},
117120
}
118121

119-
# Add activation memory if enabled
120-
if config_dict.get("enable_activation_memory", False):
121-
config_dict["act_mem"] = {
122-
"backend": "kv_cache",
123-
"config": {
124-
"memory_filename": kwargs.get(
125-
"activation_memory_filename", "activation_memory.pickle"
126-
),
127-
"extractor_llm": {
128-
"backend": "openai",
129-
"config": openai_config,
130-
},
131-
},
132-
}
122+
# Note: act_mem configuration belongs in MemCube config (get_default_cube_config),
123+
# not in MOSConfig which doesn't have an act_mem field (extra="forbid").
124+
# The enable_activation_memory flag above is sufficient for MOSConfig.
133125

134126
return MOSConfig(**config_dict)
135127

@@ -237,21 +229,33 @@ def get_default_cube_config(
237229
},
238230
}
239231

240-
# Configure activation memory if enabled
232+
# Configure activation memory if enabled.
233+
# KV cache activation memory requires a local HuggingFace/vLLM model (it
234+
# extracts internal attention KV tensors via build_kv_cache), so it cannot
235+
# work with remote API backends like OpenAI.
236+
# Only create act_mem when activation_memory_backend is explicitly provided.
241237
act_mem_config = {}
242238
if kwargs.get("enable_activation_memory", False):
243-
act_mem_config = {
244-
"backend": "kv_cache",
245-
"config": {
246-
"memory_filename": kwargs.get(
247-
"activation_memory_filename", "activation_memory.pickle"
248-
),
249-
"extractor_llm": {
250-
"backend": "openai",
251-
"config": openai_config,
239+
extractor_backend = kwargs.get("activation_memory_backend")
240+
if extractor_backend in ("huggingface", "huggingface_singleton", "vllm"):
241+
act_mem_config = {
242+
"backend": "kv_cache",
243+
"config": {
244+
"memory_filename": kwargs.get(
245+
"activation_memory_filename", "activation_memory.pickle"
246+
),
247+
"extractor_llm": {
248+
"backend": extractor_backend,
249+
"config": kwargs.get("activation_memory_llm_config", {}),
250+
},
252251
},
253-
},
254-
}
252+
}
253+
else:
254+
logger.info(
255+
"Activation memory (kv_cache) requires a local model backend "
256+
"(huggingface/vllm) via activation_memory_backend kwarg. "
257+
"Skipping act_mem in MemCube config."
258+
)
255259

256260
# Create MemCube configuration
257261
cube_config_dict = {

0 commit comments

Comments
 (0)