fix: parse model_kwargs (#1825)

src/axolotl/utils/config/models/input/v0_4_1/__init__.py

@@ -321,6 +321,8 @@ class ModelInputConfig(BaseModel):
     )
     trust_remote_code: Optional[bool] = None
 
+    model_kwargs: Optional[Dict[str, Any]] = None
+
     @field_validator("trust_remote_code")
     @classmethod
     def hint_trust_remote_code(cls, trust_remote_code):
@@ -614,6 +616,8 @@ class AxolotlInputConfig(
     flash_attn_fuse_mlp: Optional[bool] = None
     flash_optimum: Optional[bool] = None
 
+    eager_attention: Optional[bool] = None
+
     unsloth_cross_entropy_loss: Optional[bool] = None
     unsloth_lora_mlp: Optional[bool] = None
     unsloth_lora_qkv: Optional[bool] = None

src/axolotl/utils/models.py

@@ -655,19 +655,11 @@ def load_model(
                 if "device_map" in model_kwargs:
                     del model_kwargs["device_map"]
 
-            if cfg.fsdp and not cfg.adapter and cfg.local_rank != 0:
-                with init_empty_weights():
-                    model = AutoModelForCausalLM.from_pretrained(
-                        base_model,
-                        config=model_config,
-                        **model_kwargs,
-                    )
-            else:
-                model = AutoModelForCausalLM.from_pretrained(
-                    base_model,
-                    config=model_config,
-                    **model_kwargs,
-                )
+            model = AutoModelForCausalLM.from_pretrained(
+                base_model,
+                config=model_config,
+                **model_kwargs,
+            )
 
             if cfg.flash_attention and not inference:
                 from axolotl.monkeypatch.llama_attn_hijack_flash import (