Add optional start_method argument to _cli.py

src/fairchem/core/_cli.py

@@ -79,6 +79,12 @@ class DistributedInitMethod(StrEnum):
     FILE = "file"
 
 
+class StartMethod(StrEnum):
+    SPAWN = "spawn"
+    FORK = "fork"
+    FORKSERVER = "forkserver"
+
+
 @dataclass
 class SlurmConfig:
     mem_gb: int = 80
@@ -91,12 +97,14 @@ class SlurmConfig:
     account: Optional[str] = (
         None  # omegaconf in python 3.9 does not backport annotations
     )
+    additional_parameters: dict = field(default_factory=lambda: {})
 
 
 @dataclass
 class SchedulerConfig:
     mode: SchedulerType = SchedulerType.LOCAL
     distributed_init_method: DistributedInitMethod = DistributedInitMethod.TCP
+    start_method: StartMethod = StartMethod.SPAWN # this is only used when `mode=LOCAL`
     ranks_per_node: int = 1
     num_nodes: int = 1
     num_array_jobs: int = 1
@@ -444,6 +452,7 @@ def main(
             nodes=scheduler_cfg.num_nodes,
             slurm_qos=scheduler_cfg.slurm.qos,
             slurm_account=scheduler_cfg.slurm.account,
+            slurm_additional_parameters=scheduler_cfg.slurm.additional_parameters,
         )
         if scheduler_cfg.num_array_jobs == 1:
             job = executor.submit(Submitit(), cfg)
@@ -493,6 +502,7 @@ def main(
                 min_nodes=1,
                 max_nodes=1,
                 nproc_per_node=scheduler_cfg.ranks_per_node,
+                start_method=scheduler_cfg.start_method,
                 rdzv_backend="c10d",
                 max_restarts=0,
             )

src/fairchem/core/components/common/dataloader_builder.py

@@ -15,7 +15,7 @@
 
 
 def get_dataloader(
-    dataset: Dataset, batch_sampler_fn: callable, collate_fn: callable, num_workers
+    dataset: Dataset, batch_sampler_fn: callable, collate_fn: callable, num_workers, pin_memory=True,
 ) -> DataLoader:
     if gp_utils.initialized():
         num_replicas = gp_utils.get_dp_world_size()
@@ -35,7 +35,7 @@ def get_dataloader(
         dataset=dataset,
         collate_fn=collate_fn,
         num_workers=num_workers,
-        pin_memory=True,
+        pin_memory=pin_memory,
         batch_sampler=batch_sampler,
     )
     logging.info("get_dataloader::Done!")