From 0bc114d2e18c9cb148e4e2c419b49ab473e15bfd Mon Sep 17 00:00:00 2001 From: David Baker Date: Mon, 11 Mar 2024 00:50:04 +0000 Subject: [PATCH] Fix pydantic configuration for the max_memory input (#1385) [skip ci] * Fix pydantic configuration for the max_memory input * chore: lint --------- Co-authored-by: Wing Lian --- src/axolotl/utils/config/models/input/v0_4_1/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py index 11983acd9..5eab9f61e 100644 --- a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py +++ b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py @@ -512,7 +512,9 @@ class AxolotlInputConfig( neftune_noise_alpha: Optional[float] = None - max_memory: Optional[Union[int, str]] = None + max_memory: Optional[ + Dict[Union[int, Literal["cpu", "disk"]], Union[int, str]] + ] = None gpu_memory_limit: Optional[Union[int, str]] = None chat_template: Optional[Union[Literal["chatml", "inst"], ChatTemplate]] = None