fix max_memory (#3347)

Author: Jintao-Huang

swift/llm/argument/base_args/model_args.py

@@ -1,4 +1,5 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
+import ast
 import math
 import os
 from dataclasses import dataclass, field
@@ -80,6 +81,11 @@ def _init_device_map(self):
                     self.device_map[k] += local_rank
 
     def _init_max_memory(self):
+        if isinstance(self.max_memory, str):
+            try:
+                self.max_memory = ast.literal_eval(self.max_memory)
+            except Exception:
+                pass
         self.max_memory = self.parse_to_dict(self.max_memory)
         # compat mp&ddp
         _, local_rank, _, local_world_size = get_dist_setting()

swift/llm/model/register.py

@@ -486,7 +486,7 @@ def get_model_tokenizer(
         # model kwargs
         model_type: Optional[str] = None,
         quantization_config=None,
-        max_memory: Optional[List[str]] = None,
+        max_memory: Union[str, Dict[str, Any]] = None,
         attn_impl: Literal['flash_attn', 'sdpa', 'eager', None] = None,
         rope_scaling: Optional[Dict[str, Any]] = None,
         automodel_class=None,