Merge branch 'main' into add_calculator_only_output_field

FAIR-Chem · Dec 5, 2024 · 15c946b · 15c946b
2 parents 4c7c8dd + 6ab6ad7
commit 15c946b
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 4 deletions.
diff --git a/packages/env.cpu.yml b/packages/env.cpu.yml
@@ -4,7 +4,7 @@ channels:
 - defaults
 dependencies:
 - cpuonly
-- pytorch>=2.4
+- pytorch==2.4.0
 - ase
 - e3nn>=0.5
 - numpy >=1.26.0,<2.0.0

diff --git a/packages/env.gpu.yml b/packages/env.gpu.yml
@@ -5,7 +5,7 @@ channels:
 - defaults
 dependencies:
 - pytorch-cuda=12.1
-- pytorch>=2.4
+- pytorch==2.4.0
 - ase
 - e3nn>=0.5
 - numpy >=1.26.0,<2.0.0

diff --git a/packages/fairchem-core/pyproject.toml b/packages/fairchem-core/pyproject.toml
@@ -9,7 +9,7 @@ license = {text = "MIT License"}
 dynamic = ["version", "readme"]
 requires-python = ">=3.9, <3.13"
 dependencies = [
-    "torch>=2.4",
+    "torch==2.4",
     "numpy >=1.26.0, <2.0.0",
     "lmdb",
     "ase",

diff --git a/src/fairchem/core/common/distutils.py b/src/fairchem/core/common/distutils.py
@@ -80,7 +80,8 @@ def setup(config) -> None:
                     assign_device_for_local_rank(config["cpu"], config["local_rank"])
                 else:
                     # in the old code, all ranks can see all devices but need to be assigned a device equal to their local rank
-                    # this is dangerous and should be deprecated
+                    # this is dangerous and should be deprecated, however, FSDP still requires backwards compatibility with
+                    # initializing this way for now so we need to keep it
                     torch.cuda.set_device(config["local_rank"])
 
                 dist.init_process_group(
@@ -123,6 +124,11 @@ def setup(config) -> None:
         config["local_rank"] = int(os.environ.get("LOCAL_RANK"))
         if config.get("use_cuda_visibile_devices"):
             assign_device_for_local_rank(config["cpu"], config["local_rank"])
+        elif torch.cuda.is_available():
+            # in the old code, all ranks can see all devices but need to be assigned a device equal to their local rank
+            # this is dangerous and should be deprecated, however, FSDP still requires backwards compatibility with
+            # initializing this way for now so we need to keep it
+            torch.cuda.set_device(config["local_rank"])
         dist.init_process_group(
             backend=config["distributed_backend"],
             rank=int(os.environ.get("RANK")),