Fix Deepspeed loading (#950)

* add check for zero3

* freeze parameters

* fixes for deepspeed loading

* fix model parameter check

* unfrozen parameters in example mixtral and logging when unfreezing

src/axolotl/train.py

@@ -18,6 +18,7 @@ from axolotl.common.cli import TrainerCliArgs
 from axolotl.logging_config import configure_logging
 from axolotl.monkeypatch import neft_embeddings
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.freeze import freeze_parameters_except
 from axolotl.utils.models import load_model, load_tokenizer
 from axolotl.utils.trainer import setup_trainer
 
@@ -78,6 +79,9 @@ def train(
             )
     resume_from_checkpoint = cfg.resume_from_checkpoint
 
+    if cfg.unfrozen_parameters:
+        freeze_parameters_except(model, cfg.unfrozen_parameters)
+
     trainer = setup_trainer(
         cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps
     )