SLURM + FSDP2 support

pyproject.toml

@@ -76,7 +76,7 @@ nn = [
     "dask-jobqueue",
     "dask[distributed]",
     "tokenizers>0.20",
-    "torch>=2.5,<2.7",  # Temporarily pin <2.6 until someone fixes our CI with torch 2.6
+    "torch>=2.6,<2.7",  # Use torch 2.6
     "transformers>=4.46",
     "wandb",
 ]
@@ -153,7 +153,7 @@ explicit_package_bases = true
 mypy_path = "$MYPY_CONFIG_FILE_DIR/src,$MYPY_CONFIG_FILE_DIR/packages/lmi/src"
 # Specifies the OS platform for the target program, for example darwin or win32
 # (meaning OS X or Windows, respectively). The default is the current platform
-# as revealed by Python’s sys.platform variable.
+# as revealed by Python's sys.platform variable.
 platform = "linux"
 # Comma-separated list of mypy plugins.
 plugins = ["pydantic.mypy"]

src/ldp/nn/__init__.py

@@ -12,6 +12,7 @@
     ParallelTransformerHandler,
     TransformerHandler,
     TransformerHandlerConfig,
+    ParallelizationStrategy,
     collate_fn_transformer_left_pad,
     collate_fn_transformer_right_pad,
     decollate_fn_transformer_decoder,
@@ -35,6 +36,7 @@
     "TorchDType",
     "TransformerHandler",
     "TransformerHandlerConfig",
+    "ParallelizationStrategy",
     "collate_fn_transformer_left_pad",
     "collate_fn_transformer_right_pad",
     "decollate_fn_transformer_decoder",

src/ldp/nn/agent/simple_local_agent.py

@@ -15,6 +15,7 @@
 from ldp.nn.handlers.chunking import TensorChunker
 from ldp.nn.handlers.transformer_handler import (
     ParallelModeConfig,
+    ParallelizationStrategy,
     logits_to_logprobs,
 )
 from ldp.nn.lm_config import LMConfig as _LMConfig
@@ -31,7 +32,7 @@ class AgentLMConfig(_LMConfig):
 
     # distribution
     parallel_mode: ParallelModeConfig | None = None
-
+    parallel_strategy: ParallelizationStrategy = ParallelizationStrategy.ACCELERATOR
     # sampling parameters
     temperature: float = 1.0
     max_new_tokens: int = 50
@@ -80,6 +81,7 @@ def __init__(
             batch_size=self.llm_model.batch_size,
             max_wait_interval=self.llm_model.max_wait_interval,
             parallel_mode_config=self.llm_model.parallel_mode,
+            parallel_strategy=self.llm_model.parallel_strategy,
         )
 
     async def init_state(self, tools: list[Tool]) -> SimpleAgentState:

src/ldp/nn/graph/llm_call_op.py

@@ -16,6 +16,7 @@
     LMType,
     ParallelModeConfig,
     TransformerHandlerConfig,
+    ParallelizationStrategy,
     collate_fn_transformer_left_pad,
     decollate_fn_transformer_decoder,
 )
@@ -40,6 +41,7 @@ def __init__(
         batch_size: int = 1,
         max_wait_interval: float = 0.1,
         parallel_mode_config: ParallelModeConfig | None = None,
+        parallel_strategy: ParallelizationStrategy = ParallelizationStrategy.ACCELERATOR,
     ) -> None:
         super().__init__()
 
@@ -51,6 +53,7 @@ def __init__(
             batch_size=batch_size,
             max_wait_interval=max_wait_interval,
             parallel_mode_config=parallel_mode_config,
+            parallel_strategy=parallel_strategy,
             # constant configuration
             lm_type=LMType.GENERATION,
             module_call_fn=AsyncTransformerInterface.model_generate,