fixing tests and imports

Author: SalmanMohammadi

src/axolotl/utils/schemas/validation.py

@@ -759,7 +759,7 @@ def check_muon_deepspeed_fsdp(cls, data):
             if data.get("fsdp") or data.get("fsdp_config"):
                 fsdp_version = data.get("fsdp_version")
                 if fsdp_version is None:
-                    fsdp_version = fsdp_config.get("fsdp_version", 1)
+                    fsdp_version = data.get("fsdp_config", {}).get("fsdp_version", 1)
                 if str(fsdp_version) != "2":
                     raise ValueError(
                         "Muon optimizer is only compatible with FSDP2. Set fsdp_version: 2 to use Muon with FSDP."

tests/core/test_builders.py

@@ -474,10 +474,7 @@ def test_custom_optimizer_cls_and_kwargs(
 
             assert trainer.optimizer_cls_and_kwargs is not None
 
-            from axolotl.contribs.mit.muon import (
-                Muon,
-                MuonOptimizerFactory,
-            )
+            from axolotl.contribs.mit.muon import MuonOptimizerFactory
 
             optimizer_cls, optimizer_kwargs = trainer.optimizer_cls_and_kwargs
             assert optimizer_cls is MuonOptimizerFactory
@@ -556,10 +553,7 @@ def test_custom_optimizer_cls_and_kwargs(
 
         assert trainer.optimizer_cls_and_kwargs is not None
 
-        from axolotl.contribs.mit.muon import (
-            Muon,
-            MuonOptimizerFactory,
-        )
+        from axolotl.contribs.mit.muon import MuonOptimizerFactory
 
         optimizer_cls, optimizer_kwargs = trainer.optimizer_cls_and_kwargs
         assert optimizer_cls is MuonOptimizerFactory