hotfix for capabilities loading (#1331)

src/axolotl/cli/__init__.py

@@ -30,7 +30,6 @@ from axolotl.common.cli import TrainerCliArgs, load_model_and_tokenizer
 from axolotl.logging_config import configure_logging
 from axolotl.train import TrainDatasetMeta
 from axolotl.utils.config import (
-    GPUCapabilities,
     normalize_cfg_datasets,
     normalize_config,
     validate_config,
@@ -350,14 +349,15 @@ def load_cfg(config: Union[str, Path] = Path("examples/"), **kwargs):
     except:  # pylint: disable=bare-except # noqa: E722
         gpu_version = None
 
-    capabilities = GPUCapabilities(
-        bf16=is_torch_bf16_gpu_available(),
-        n_gpu=os.environ.get("WORLD_SIZE", 1),
-        compute_capability=gpu_version,
+    cfg = validate_config(
+        cfg,
+        capabilities={
+            "bf16": is_torch_bf16_gpu_available(),
+            "n_gpu": os.environ.get("WORLD_SIZE", 1),
+            "compute_capability": gpu_version,
+        },
     )
 
-    cfg = validate_config(cfg, capabilities=capabilities)
-
     prepare_optim_env(cfg)
 
     normalize_config(cfg)

src/axolotl/utils/config/__init__.py

@@ -13,7 +13,6 @@ from axolotl.utils.config.models.input.v0_4_1 import (
     AxolotlConfigWCapabilities,
     AxolotlInputConfig,
 )
-from axolotl.utils.config.models.internals import GPUCapabilities
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.models import load_model_config
 
@@ -197,7 +196,7 @@ def normalize_cfg_datasets(cfg):
                     cfg.datasets[idx].conversation = "chatml"
 
 
-def validate_config(cfg: DictDefault, capabilities: Optional[GPUCapabilities] = None):
+def validate_config(cfg: DictDefault, capabilities: Optional[dict] = None):
     if capabilities:
         return DictDefault(
             dict(AxolotlConfigWCapabilities(**cfg.to_dict(), capabilities=capabilities))