add torch_compile_mode options (#1763) [skip ci]

* add torch_compile_mode options * make sure n_gpu is an int
2024-07-17 15:38:07 -04:00
parent 976f85195a
commit 8619b2d855
3 changed files with 17 additions and 1 deletions
--- a/src/axolotl/cli/init.py
+++ b/src/axolotl/cli/init.py
@@ -375,7 +375,7 @@ def load_cfg(config: Union[str, Path] = Path("examples/"), **kwargs):
        cfg,
        capabilities={
            "bf16": is_torch_bf16_gpu_available(),
-            "n_gpu": os.environ.get("WORLD_SIZE", 1),
+            "n_gpu": int(os.environ.get("WORLD_SIZE", 1)),
            "compute_capability": gpu_version,
        },
    )