feat: optional multi-GPU (DDP) finetuning

agent/config/defaults_finetune.json

@@ -37,5 +37,5 @@
 
   "_about_mtl": "Per-target task-specific FFN heads are exposed via ffn_num_task_specific_layers (default 0 == off) and ffn_task_specific_hidden_size (default null; required when ffn_num_task_specific_layers > 0). When >0, each target column gets its own N-layer FFN stacked on the shared trunk; useful for heterogeneous multi-target finetunes (e.g. solubility + permeability + metabolic stability together). Override via --ffn-num-task-specific-layers N --ffn-task-specific-hidden-size H.",
 
-  "_about_gpu_selection": "GPU selection is auto-detected at runtime, not a default here. Finetune defaults to GPU 0; override with --gpus 0,2 etc."
+  "_about_gpu_selection": "GPU selection is auto-detected at runtime, not a default here. Finetune defaults to GPU 0. Use --gpus <id> to pick a single device for single-GPU finetune; use --num-gpus N for multi-GPU data-parallel (DDP)."
 }

agent/scripts/run_finetune_local.py

@@ -19,9 +19,10 @@
 User-supplied arch flags are not exposed; the runner refuses to override what
 the ckpt dictates so the FFN heads attach to a consistent encoder.
 
-Single-GPU is the default for finetune (per agent/README's hardware table —
-multi-GPU finetune is not currently supported and main.py's finetune path
-doesn't DDP). `--gpus N` picks a specific device id; defaults to GPU 0.
+Single-GPU is the default for finetune. `--gpus N` picks a specific device id
+(defaults to GPU 0). For data-parallel multi-GPU finetuning, pass `--num-gpus N`
+(N>1): the runner sets `WORLD_SIZE=N` and `main.py finetune` spawns one process
+per GPU, and `--batch-size` is interpreted per-GPU.
 
 CLI
 ---
@@ -189,8 +190,16 @@ def _validate_mtl_consistency(applied: dict[str, dict[str, Any]]) -> None:
 def _build_argv(
     *, gpu: int, out_dir: Path, manifest: dict[str, Any], ckpt: Path,
     arch: dict[str, Any], applied: dict[str, dict[str, Any]],
+    num_gpus: int = 1,
 ) -> list[str]:
-    """Constructs the full `main.py finetune` argv as a list of strings."""
+    """Constructs the full finetune argv as a list of strings.
+
+    Single-process and multi-GPU (DDP) finetune share one entrypoint,
+    `main.py finetune ...`. DDP is selected at runtime by the WORLD_SIZE env var
+    (this runner sets it to num_gpus when num_gpus > 1; main.py then spawns one
+    process per GPU). The argv is therefore identical for both; num_gpus is kept
+    only to document that --batch_size is per-GPU under DDP.
+    """
     outputs = manifest["outputs"]
     method = manifest["split_method"]
 
@@ -273,8 +282,11 @@ def _build_argv(
     if applied.get("show_individual_scores", {}).get("value"):
         argv += ["--show_individual_scores"]
 
-    # GPU + save_dir
-    argv += ["--gpu", str(gpu)]
+    # GPU + save_dir. Under DDP (num_gpus>1) gpu is None: main.py finetune pins
+    # each rank to its own device, and run_training ignores args.gpu when
+    # distributed, so no --gpu flag is passed.
+    if gpu is not None:
+        argv += ["--gpu", str(gpu)]
     argv += ["--save_dir", str(out_dir / "ckpt")]
 
     return argv
@@ -309,25 +321,33 @@ def run(args: argparse.Namespace) -> dict[str, Any]:
     arch = _arch_from_validator(validator_out)
     model_type = validator_out.get("model_type")
 
-    # 3. GPU selection (single-GPU only).
-    gpu = resolve_single_gpu(args.gpus, workflow="finetune")
+    # 3. GPU selection.
+    num_gpus = max(1, int(getattr(args, "num_gpus", 1) or 1))
+    distributed = num_gpus > 1
+    if distributed:
+        # Data-parallel DDP: main.py finetune uses GPUs 0..num_gpus-1 (one rank
+        # each) via WORLD_SIZE; a single-device pin does not apply.
+        gpu = None
+    else:
+        gpu = resolve_single_gpu(args.gpus, workflow="finetune")
 
     # 4. Apply defaults + collect args_applied.
     applied = _apply_defaults(args, defaults)
     # MTL FFN consistency
     _validate_mtl_consistency(applied)
 
-    # 5. Build the main.py finetune argv.
+    # 5. Build the finetune argv (always main.py finetune; DDP selected via WORLD_SIZE).
     argv = _build_argv(
         gpu=gpu, out_dir=out_dir, manifest=manifest, ckpt=ckpt,
-        arch=arch, applied=applied,
+        arch=arch, applied=applied, num_gpus=num_gpus,
     )
 
     # 6. Build the run.json manifest.
     commit, dirty = git_commit_with_env_override(REPO_ROOT)
     image_tag = os.environ.get("KERMT_IMAGE", "kermt:latest")
     image_digest = docker_image_digest(image_tag)
-    cmd_replay = format_cmd_replay(argv, env={"CUDA_VISIBLE_DEVICES": gpu})
+    replay_env = {"WORLD_SIZE": str(num_gpus)} if distributed else {"CUDA_VISIBLE_DEVICES": gpu}
+    cmd_replay = format_cmd_replay(argv, env=replay_env)
     targets = manifest.get("targets")
     run_manifest: dict[str, Any] = {
         "workflow": "finetune",
@@ -344,6 +364,8 @@ def run(args: argparse.Namespace) -> dict[str, Any]:
         },
         "model_type": model_type,
         "gpu": gpu,
+        "num_gpus": num_gpus,
+        "distributed": distributed,
         "args_applied": applied,
         "arch": arch,
         "save_dir": str(out_dir / "ckpt"),
@@ -362,10 +384,14 @@ def run(args: argparse.Namespace) -> dict[str, Any]:
         return run_manifest
 
     env = os.environ.copy()
-    env["CUDA_VISIBLE_DEVICES"] = str(gpu)
+    if distributed:
+        # DDP: main.py finetune reads WORLD_SIZE and spawns one process per GPU.
+        # Do not pin CUDA_VISIBLE_DEVICES to a single device.
+        env["WORLD_SIZE"] = str(num_gpus)
+    else:
+        env["CUDA_VISIBLE_DEVICES"] = str(gpu)
     # main.py enables strict deterministic algorithms via
-    # `torch.use_deterministic_algorithms(True)` (kermt/main.py:23); CuBLAS
-    # then requires this env var to be set.
+    # `torch.use_deterministic_algorithms(True)`; CuBLAS then requires this env var.
     env.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
 
     log_file = out_dir / "logs" / "finetune.log"
@@ -389,8 +415,13 @@ def main(argv: list[str] | None = None) -> int:
     p.add_argument("--ckpt-validator-out", default=None,
                    help="Optional cached check_checkpoint.py JSON; computed if absent.")
     p.add_argument("--gpus", default=None,
-                   help="Single GPU id (default 0). Finetune is single-GPU only; "
-                        "passing '0,1' is rejected with a clear error.")
+                   help="Single GPU id for single-process finetune (default 0). "
+                        "Ignored when --num-gpus > 1 (DDP uses ranks 0..N-1).")
+    p.add_argument("--num-gpus", type=int, default=1,
+                   help="Number of GPUs for data-parallel (DDP) finetune. Default 1 "
+                        "(single-process main.py finetune, unchanged). N>1 runs "
+                        "main.py finetune with WORLD_SIZE=N (one process per GPU); "
+                        "--batch-size is per-GPU.")
     p.add_argument("--dry-run", action="store_true",
                    help="Write run.json + print the command without executing.")
 

agent/skills/kermt-finetune/SKILL.md

@@ -19,9 +19,10 @@ launch the runner detached, return a run directory + container name.
 
 ## Hardware requirements
 
-- **GPUs**: 1 (single-GPU). Finetune does not DDP; if you have multiple
-  GPUs visible, pass `--gpus 0` (or whichever id) to select one. Multi-GPU
-  finetune is not currently supported.
+- **GPUs**: 1 by default (single-GPU); pass `--gpus 0` (or whichever id) to
+  select one. For faster training on a multi-GPU host, pass `--num-gpus N`
+  (N>1) to run data-parallel DDP across N GPUs — `--batch-size` is then
+  per-GPU (effective global batch = batch_size × N).
 - **VRAM**: ≥ 8 GB for the default `batch_size 32` configuration. Lower VRAM
   works at smaller batch sizes — pass `--batch-size N` to override.
 - **Disk**: a few GB per run (checkpoint + features + logs).
@@ -84,7 +85,11 @@ Optional:
   finetunes). Both must be set together when N > 0.
 - `--ensemble-size N` / `--num-folds N` — multi-model / k-fold CV. Default 1
   each.
-- `--gpus 0` — single GPU id (default 0). Passing more than one is rejected.
+- `--gpus 0` — single GPU id for single-process finetune (default 0). Ignored
+  when `--num-gpus > 1`.
+- `--num-gpus N` — number of GPUs for data-parallel DDP finetune. Default 1
+  (single-process, unchanged). N>1 runs `main.py finetune` with `WORLD_SIZE=N`
+  (one process per GPU); `--batch-size` is per-GPU.
 - `--from-prepare <dir>` — skip the prepare step and reuse an existing
   `prepare_data.json` in `<dir>`. Useful when iterating on hyperparameters.
 
@@ -224,6 +229,7 @@ helper bind-mounts them at known container paths.
             --dataset-type <type> \\
             --out /runs \\
             [--gpus 0] \\
+            [--num-gpus N] \\
             [--epochs N --batch-size N --init-lr F ...] \\
             [--ffn-num-task-specific-layers N --ffn-task-specific-hidden-size H]"
    ```
@@ -279,8 +285,8 @@ helper bind-mounts them at known container paths.
   step (typically clean_smiles or save_features). Fix and re-run.
 - `ffn_num_task_specific_layers=N>0 but ffn_task_specific_hidden_size is unset`
   → MTL heads need an explicit hidden size. Pass `--ffn-task-specific-hidden-size H`.
-- `finetune is single-GPU` → multi-GPU finetune is not currently supported;
-  pick a single id.
+- `finetune is single-GPU` (from `--gpus 0,1`) → `--gpus` selects one device
+  for single-process finetune. For multi-GPU, use `--num-gpus N` (DDP) instead.
 
 ## Replayability
 

agent/tests/test_run_finetune_local.py

@@ -457,6 +457,28 @@ def test_gpu_id_propagates(tmp_path: Path) -> None:
     assert argv[argv.index("--gpu") + 1] == "2"
 
 
+def test_multi_gpu_uses_main_entrypoint(tmp_path: Path) -> None:
+    """DDP finetune (--num-gpus > 1) uses the same `main.py finetune` entrypoint
+    (the finetune_ddp.py launcher was removed); DDP is selected via WORLD_SIZE."""
+    ckpt, prep, val = _setup(tmp_path)
+    code, m = _run(
+        "--ckpt", str(ckpt), "--prepare-manifest", str(prep),
+        "--ckpt-validator-out", str(val), "--out", str(tmp_path / "run"),
+        "--dataset-type", "regression",
+        "--num-gpus", "2", "--dry-run",
+    )
+    assert code == 0, m
+    manifest = m["manifest"]
+    assert manifest["num_gpus"] == 2
+    assert manifest["gpu"] is None  # DDP: no single-device pin
+    argv = manifest["argv"]
+    assert "main.py" in argv[2]
+    assert argv[3] == "finetune"
+    assert all("finetune_ddp" not in str(a) for a in argv)
+    # DDP is selected at runtime via WORLD_SIZE in the replay command.
+    assert "WORLD_SIZE" in manifest["cmd_replay"]
+
+
 # ---------------------------------------------------------------------------
 # Manifest schema
 # ---------------------------------------------------------------------------

kermt/util/ddp_utils.py

@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Shared DistributedDataParallel (DDP) helpers for KERMT.
+
+Used by the pretraining and finetuning DDP launchers so both share a single,
+proven single-node DDP bootstrap. The logic here matches what pretraining has
+used in production (localhost rendezvous, NCCL backend, per-process GPU pinning
+by rank, and topology-aware NCCL P2P configuration).
+"""
+import os
+import subprocess
+
+import torch
+from torch.distributed import init_process_group
+
+
+def configure_nccl_for_topology():
+    """
+    Auto-configure NCCL settings based on GPU topology.
+    This handles cases where P2P (peer-to-peer) GPU communication is not available.
+    Must be called BEFORE spawning processes (in main process).
+    """
+    # Check if user has already set NCCL settings (don't override)
+    if "NCCL_P2P_DISABLE" in os.environ:
+        print(f"[INFO] Using user-provided NCCL settings: NCCL_P2P_DISABLE={os.environ['NCCL_P2P_DISABLE']}")
+        return
+
+    # Try to detect GPU topology
+    try:
+        result = subprocess.run(['nvidia-smi', 'topo', '-m'],
+                                capture_output=True, text=True, timeout=5)
+        topo_output = result.stdout
+
+        # Check for poor GPU connectivity (SYS or NODE topology)
+        # These topologies typically don't support P2P well
+        if 'SYS' in topo_output or 'NODE' in topo_output:
+            print("[INFO] Detected cross-NUMA or system-level GPU topology (SYS/NODE).")
+            print("[INFO] Disabling P2P for stability. This is normal for multi-socket systems.")
+            os.environ["NCCL_P2P_DISABLE"] = "1"
+            os.environ["NCCL_IB_DISABLE"] = "1"
+            os.environ["NCCL_SHM_DISABLE"] = "0"
+        else:
+            print("[INFO] GPU topology appears to support P2P. Enabling P2P communication.")
+    except Exception as e:
+        # If detection fails, use safe defaults (disable P2P)
+        print(f"[WARNING] Could not detect GPU topology: {e}")
+        print("[INFO] Using safe default: P2P disabled. Set NCCL_P2P_DISABLE=0 to enable if your system supports it.")
+        os.environ["NCCL_P2P_DISABLE"] = "1"
+        os.environ["NCCL_IB_DISABLE"] = "1"
+        os.environ["NCCL_SHM_DISABLE"] = "0"
+
+
+def ddp_setup(rank: int, world_size: int):
+    """
+    Initialize the process group for single-node DDP and pin this process to its GPU.
+
+    Args:
+        rank: Unique identifier of each process (also the GPU index it is pinned to).
+        world_size: Total number of processes.
+    """
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12355"
+    torch.cuda.set_device(rank)
+    init_process_group(backend="nccl", rank=rank, world_size=world_size)

kermt/util/utils.py

@@ -63,6 +63,28 @@
 from kermt.util.scheduler import NoamLR
 
 
+def seed_rngs(seed: int) -> None:
+    """Seed the torch (CPU + CUDA), numpy, and python RNGs."""
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+
+def setup_determinism(seed: int) -> None:
+    """Seed all RNGs and enable deterministic algorithms.
+
+    Shared by the finetune / HPO entry points (main.py -- single-process and DDP
+    finetune -- and main_hpo.py) so their determinism setup stays in lockstep. Also sets
+    CUBLAS_WORKSPACE_CONFIG, which torch.use_deterministic_algorithms requires
+    for cuBLAS on CUDA >= 10.2.
+    """
+    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+    seed_rngs(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.use_deterministic_algorithms(mode=True)
+
+
 def get_model_args():
     """
     Get model structure related parameters.
@@ -673,13 +695,18 @@ def build_optimizer(model: nn.Module, args: Namespace):
     return optimizer
 
 
-def build_lr_scheduler(optimizer, args: Namespace, total_epochs: List[int] = None):
+def build_lr_scheduler(optimizer, args: Namespace, total_epochs: List[int] = None,
+                       world_size: int = 1):
     """
     Builds a learning rate scheduler.
 
     :param optimizer: The Optimizer whose learning rate will be scheduled.
     :param args: Arguments.
     :param total_epochs: The total number of epochs for which the model will be task.
+    :param world_size: number of DDP processes. Under DDP each rank sees
+        1/world_size of the data (via DistributedSampler), so the number of
+        optimizer steps per epoch is divided by world_size (matches the effective
+        global batch of batch_size * world_size, as in pretrain_ddp.py).
     :return: An initialized learning rate scheduler.
     """
 
@@ -688,12 +715,12 @@ def build_lr_scheduler(optimizer, args: Namespace, total_epochs: List[int] = Non
     # so we only have task params (1 group) with full LR (fine_tune_coff=1.0)
     # When fine_tune_coff > 0, we have 2 groups: encoder (index 0) and task (index 1)
     scheduler_fine_tune_coff = 1.0 if args.fine_tune_coff == 0 else args.fine_tune_coff
-    
+
     return NoamLR(
         optimizer=optimizer,
         warmup_epochs=args.warmup_epochs,
         total_epochs=args.epochs,
-        steps_per_epoch=args.train_data_size // args.batch_size,
+        steps_per_epoch=args.train_data_size // (args.batch_size * world_size),
         init_lr=args.init_lr,
         max_lr=args.max_lr,
         final_lr=args.final_lr,