feat(config): inline DeepSpeed config into main YAML (#36)

Neonkraft · web-flow · commit d0d0bb90aaa8 · 2026-07-02T13:33:51.000+02:00
* feat(config): inline DeepSpeed config into main YAML, drop separate config file

* fix(tests): update test_config.py for inline deepspeed config

* fix(config): type deepspeed field as dict[str, Any] | None

* fix(methods): normalize empty deepspeed dict to None

* fix(slurm): gate DeepSpeed accelerate flags on the deepspeed config

* fix(config): reject deprecated deepspeed.config_path style

* docs: note how to switch DeepSpeed ZeRO stage 2 to 3
diff --git a/README.md b/README.md
@@ -302,7 +302,8 @@ You must specify exactly one determining factor for training duration in the `tr
 
 ### 4. Infrastructure & Compute
 
-- **DeepSpeed**: configured via `deepspeed.config_path` (e.g., `configs/deepspeed/zero3.yaml`)
+- **DeepSpeed**: configured inline under the top-level `deepspeed:` key (see the reference config below); set `deepspeed: null` to disable DeepSpeed entirely.
+  To switch from ZeRO stage 2 to stage 3, bump `zero_optimization.stage` to `3` and add the `stage3_*` tuning keys — see `configs/deepspeed/zero3.yaml` for a full example.
 - **Accelerate flags**: the `accelerate` section in the YAML mirrors the CLI flags required for multi-node setups (`mixed_precision`, `dynamo_backend`, `rdzv_backend`, etc.).
   These are used by the SLURM launcher to generate the correct job script.
 - **Self-healing**: the SLURM launcher (`src/post_training/slurm/`) supports auto-requeueing.
@@ -485,8 +486,28 @@ data:
       transform: null                        # null = already conversational
 
 # -- DeepSpeed ---------------------------------------------------------------
+# Set to null to disable DeepSpeed entirely.
+# To switch from ZeRO stage 2 to stage 3, bump zero_optimization.stage to 3 and
+# add the stage3_* tuning keys — see configs/deepspeed/zero3.yaml for a full example.
 deepspeed:
-  config_path: "configs/deepspeed/zero2.yaml"
+  bf16:
+    enabled: true
+  zero_optimization:
+    stage: 2
+    overlap_comm: true
+    contiguous_gradients: true
+    reduce_scatter: true
+  gradient_clipping: 1.0
+  train_micro_batch_size_per_gpu: "auto"
+  gradient_accumulation_steps: "auto"
+  train_batch_size: "auto"
+  optimizer:
+    type: AdamW
+    params:
+      lr: "auto"
+      betas: "auto"
+      eps: "auto"
+      weight_decay: "auto"
 
 # -- Accelerate launch flags (explicit multi-node control) -------------------
 accelerate:
diff --git a/configs/trl/dpo.yaml b/configs/trl/dpo.yaml
@@ -73,8 +73,28 @@ data:
       transform: null                        # null = already "chosen" / "rejected" columns
 
 # -- DeepSpeed ---------------------------------------------------------------
+# Set to null to disable DeepSpeed entirely.
+# To switch from ZeRO stage 2 to stage 3, bump zero_optimization.stage to 3 and
+# add the stage3_* tuning keys — see configs/deepspeed/zero3.yaml for a full example.
 deepspeed:
-  config_path: "configs/deepspeed/zero2.yaml"
+  bf16:
+    enabled: true
+  zero_optimization:
+    stage: 2
+    overlap_comm: true
+    contiguous_gradients: true
+    reduce_scatter: true
+  gradient_clipping: 1.0
+  train_micro_batch_size_per_gpu: "auto"
+  gradient_accumulation_steps: "auto"
+  train_batch_size: "auto"
+  optimizer:
+    type: AdamW
+    params:
+      lr: "auto"
+      betas: "auto"
+      eps: "auto"
+      weight_decay: "auto"
 
 # -- Accelerate launch flags -------------------------------------------------
 accelerate:
diff --git a/configs/trl/sft.yaml b/configs/trl/sft.yaml
@@ -80,8 +80,28 @@ data:
       transform: null                        # null = already conversational
 
 # -- DeepSpeed ---------------------------------------------------------------
+# Set to null to disable DeepSpeed entirely.
+# To switch from ZeRO stage 2 to stage 3, bump zero_optimization.stage to 3 and
+# add the stage3_* tuning keys — see configs/deepspeed/zero3.yaml for a full example.
 deepspeed:
-  config_path: "configs/deepspeed/zero2.yaml"
+  bf16:
+    enabled: true
+  zero_optimization:
+    stage: 2
+    overlap_comm: true
+    contiguous_gradients: true
+    reduce_scatter: true
+  gradient_clipping: 1.0
+  train_micro_batch_size_per_gpu: "auto"
+  gradient_accumulation_steps: "auto"
+  train_batch_size: "auto"
+  optimizer:
+    type: AdamW
+    params:
+      lr: "auto"
+      betas: "auto"
+      eps: "auto"
+      weight_decay: "auto"
 
 # -- Accelerate launch flags (explicit multi-node control) -------------------
 accelerate:
diff --git a/src/post_training/backend.py b/src/post_training/backend.py
@@ -82,6 +82,13 @@ def validate(self, config: PostTrainingConfig) -> None:
                 f"Supported methods: {', '.join(_SUPPORTED_METHODS)}"
             )
 
+        if isinstance(config.deepspeed, dict) and "config_path" in config.deepspeed:
+            raise ValueError(
+                "deepspeed.config_path is no longer supported. Inline the DeepSpeed "
+                "JSON/YAML directly under the `deepspeed:` key (see configs/trl/sft.yaml "
+                "for an example), or set `deepspeed: null` to disable DeepSpeed."
+            )
+
         # Container validation (only when container.image is set)
         if config.container is not None and config.container.image:
             if not config.container.bind_mounts:
diff --git a/src/post_training/config.py b/src/post_training/config.py
@@ -11,7 +11,6 @@
 from pathlib import Path
 from typing import Any
 
-import yaml
 from omegaconf import MISSING, DictConfig, OmegaConf
 
 logger = logging.getLogger(__name__)
@@ -151,18 +150,13 @@ class DataConfig:
     datasets: list[DatasetEntry] = field(default_factory=list)
 
 
-@dataclass
-class DeepSpeedConfig:
-    """Pointer to the DeepSpeed YAML config file.  Set to ``null`` to disable."""
-
-    config_path: str | None = "configs/deepspeed/zero2.yaml"
-
-
 @dataclass
 class AccelerateConfig:
     """Flags forwarded to ``accelerate launch`` for explicit multi-node control."""
 
     mixed_precision: str = "bf16"
+    # Only takes effect when the top-level `deepspeed:` config is also set;
+    # `deepspeed: null` disables DeepSpeed at launch regardless of this flag.
     use_deepspeed: bool = True
     deepspeed_multinode_launcher: str = "standard"
     same_network: bool = True
@@ -252,7 +246,8 @@ class PostTrainingConfig:
     dpo: DPOMethodConfig = field(default_factory=DPOMethodConfig)
 
     # Infrastructure.
-    deepspeed: DeepSpeedConfig = field(default_factory=DeepSpeedConfig)
+    # Inline DeepSpeed config dict. Set to null to disable DeepSpeed entirely.
+    deepspeed: dict[str, Any] | None = None
     accelerate: AccelerateConfig = field(default_factory=AccelerateConfig)
     logging: LoggingConfig = field(default_factory=LoggingConfig)
     slurm: SlurmConfig = field(default_factory=SlurmConfig)
@@ -354,12 +349,3 @@ def resolve_gradient_accumulation_steps(self, world_size: int) -> int:
                 f"* world_size ({world_size}). Got gradient_accumulation_steps={gas}."
             )
         return int(gas)
-
-    def load_deepspeed_config(self) -> dict[str, Any]:
-        """Load the DeepSpeed YAML config and return it as a plain dict."""
-        ds_path = Path(self.deepspeed.config_path)
-        if not ds_path.is_absolute():
-            # Resolve relative to the project root (cwd).
-            ds_path = Path.cwd() / ds_path
-        with open(ds_path) as f:
-            return yaml.safe_load(f)
diff --git a/src/post_training/methods/common.py b/src/post_training/methods/common.py
@@ -75,7 +75,9 @@ def build_common_training_kwargs(
     logger.info("world_size=%d, gradient_accumulation_steps=%d", world_size, grad_accum)
 
     t = config.training
-    ds_config = config.load_deepspeed_config() if config.deepspeed.config_path else None
+    # Normalize a falsy config (e.g. `{}`) to None so it isn't forwarded to
+    # TrainingArguments as if it were an enabled DeepSpeed config.
+    ds_config = config.deepspeed or None
 
     os.environ.setdefault("TENSORBOARD_LOGGING_DIR", str(run_dir / "logs"))
 
diff --git a/src/post_training/slurm/job.sh.jinja b/src/post_training/slurm/job.sh.jinja
@@ -68,8 +68,10 @@ export LAUNCHER="accelerate launch \
     --main_process_port $MASTER_PORT \
     --mixed_precision {{ mixed_precision }} \
     --dynamo_backend {{ dynamo_backend }} \
-{% if use_deepspeed %}    --use_deepspeed \{% endif %}
+{% if use_deepspeed -%}
+    --use_deepspeed \
     --deepspeed_multinode_launcher {{ deepspeed_multinode_launcher }} \
+{% endif -%}
 {% if same_network %}    --same_network \{% endif %}
     --rdzv_backend {{ rdzv_backend }} \
     --max_restarts 0 \
diff --git a/src/post_training/slurm/job_trl_container.sh.jinja b/src/post_training/slurm/job_trl_container.sh.jinja
@@ -107,8 +107,10 @@ srun --export=ALL --wait=60 --kill-on-bad-exit=1 \
       --main_process_port $MASTER_PORT \
       --mixed_precision {{ mixed_precision }} \
       --dynamo_backend {{ dynamo_backend }} \
-{% if use_deepspeed %}      --use_deepspeed \
-{% endif %}      --deepspeed_multinode_launcher {{ deepspeed_multinode_launcher }} \
+{% if use_deepspeed -%}
+      --use_deepspeed \
+      --deepspeed_multinode_launcher {{ deepspeed_multinode_launcher }} \
+{% endif -%}
 {% if same_network %}      --same_network \
 {% endif %}      --rdzv_backend {{ rdzv_backend }} \
       --max_restarts 0 \
diff --git a/src/post_training/slurm/launcher.py b/src/post_training/slurm/launcher.py
@@ -77,7 +77,7 @@ def render_trl_slurm_script(
         # Accelerate flags
         mixed_precision=config.accelerate.mixed_precision,
         dynamo_backend=config.accelerate.dynamo_backend,
-        use_deepspeed=config.accelerate.use_deepspeed,
+        use_deepspeed=config.accelerate.use_deepspeed and bool(config.deepspeed),
         deepspeed_multinode_launcher=config.accelerate.deepspeed_multinode_launcher,
         same_network=config.accelerate.same_network,
         rdzv_backend=config.accelerate.rdzv_backend,
@@ -134,7 +134,7 @@ def render_trl_container_slurm_script(
         # Accelerate flags
         mixed_precision=config.accelerate.mixed_precision,
         dynamo_backend=config.accelerate.dynamo_backend,
-        use_deepspeed=config.accelerate.use_deepspeed,
+        use_deepspeed=config.accelerate.use_deepspeed and bool(config.deepspeed),
         deepspeed_multinode_launcher=config.accelerate.deepspeed_multinode_launcher,
         same_network=config.accelerate.same_network,
         rdzv_backend=config.accelerate.rdzv_backend,
diff --git a/src/post_training/utils/guardrails.py b/src/post_training/utils/guardrails.py
@@ -14,8 +14,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-import yaml
-
 if TYPE_CHECKING:
     from post_training.config import PostTrainingConfig
 
@@ -70,19 +68,12 @@ def _row(label: str, value: str, warn: bool = False) -> None:
 
 
 def _deepspeed_summary(config: PostTrainingConfig) -> str:
-    """Return a short description of the DeepSpeed config, e.g. 'zero2.yaml (ZeRO stage 2)'."""
-    ds_path = config.deepspeed.config_path
-    if not ds_path:
+    """Return a short description of the inline DeepSpeed config, e.g. 'ZeRO stage 2'."""
+    ds = config.deepspeed
+    if not ds:
         return "disabled"
-    path = Path(ds_path)
-    try:
-        resolved = path if path.is_absolute() else Path.cwd() / path
-        with open(resolved) as fh:
-            ds_cfg = yaml.safe_load(fh)
-        stage = ds_cfg.get("zero_optimization", {}).get("stage", "?")
-        return f"{path.name}  (ZeRO stage {stage})"
-    except Exception:
-        return str(ds_path)
+    stage = ds.get("zero_optimization", {}).get("stage", "?") if isinstance(ds, dict) else "?"
+    return f"ZeRO stage {stage}"
 
 
 # ---------------------------------------------------------------------------
diff --git a/tests/test_config.py b/tests/test_config.py
@@ -1,5 +1,6 @@
 """Tests for nullable config fields."""
 
+import pytest
 import yaml
 
 from post_training.config import PostTrainingConfig
@@ -22,7 +23,7 @@ def test_nullable_container_and_training_kwargs_load(tmp_path, monkeypatch):
                     "lr_scheduler_kwargs": None,
                     "gradient_checkpointing_kwargs": None,
                 },
-                "deepspeed": {"config_path": None},
+                "deepspeed": None,
                 "data": {
                     "datasets": [
                         {
@@ -42,3 +43,67 @@ def test_nullable_container_and_training_kwargs_load(tmp_path, monkeypatch):
     assert config.container is None
     assert kwargs["lr_scheduler_kwargs"] is None
     assert kwargs["gradient_checkpointing_kwargs"] is None
+    assert kwargs["deepspeed"] is None
+
+
+def test_deepspeed_empty_dict_normalized_to_none(tmp_path, monkeypatch):
+    monkeypatch.setenv("WORLD_SIZE", "1")
+    config_path = tmp_path / "config.yaml"
+    config_path.write_text(
+        yaml.safe_dump(
+            {
+                "method": "sft",
+                "backend": "trl",
+                "training": {
+                    "max_steps": 1,
+                    "effective_batch_size": 1,
+                    "per_device_train_batch_size": 1,
+                },
+                "deepspeed": {},
+                "data": {
+                    "datasets": [
+                        {
+                            "name": "dummy",
+                            "path": "dummy/path",
+                            "weight": 1.0,
+                        }
+                    ]
+                },
+            }
+        )
+    )
+
+    config = PostTrainingConfig.load(config_path)
+    kwargs = build_common_training_kwargs(config, tmp_path)
+
+    assert kwargs["deepspeed"] is None
+
+
+def test_deepspeed_old_style_config_path_rejected(tmp_path):
+    config_path = tmp_path / "config.yaml"
+    config_path.write_text(
+        yaml.safe_dump(
+            {
+                "method": "sft",
+                "backend": "trl",
+                "training": {
+                    "max_steps": 1,
+                    "effective_batch_size": 1,
+                    "per_device_train_batch_size": 1,
+                },
+                "deepspeed": {"config_path": "configs/deepspeed/zero2.yaml"},
+                "data": {
+                    "datasets": [
+                        {
+                            "name": "dummy",
+                            "path": "dummy/path",
+                            "weight": 1.0,
+                        }
+                    ]
+                },
+            }
+        )
+    )
+
+    with pytest.raises(ValueError, match="deepspeed.config_path is no longer supported"):
+        PostTrainingConfig.load(config_path)
diff --git a/tests/test_slurm_render.py b/tests/test_slurm_render.py
