fsdp fft loading on meta device

2024-08-11 22:18:04 -04:00
1 changed files with 13 additions and 5 deletions
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -655,6 +655,14 @@ 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,