feat: optional multi-GPU (DDP) finetuning

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 launches `finetune_ddp.py` with `WORLD_SIZE=N` instead of
+`main.py finetune`, and `--batch-size` is interpreted per-GPU.
 
 CLI
 ---
@@ -70,6 +71,7 @@
 DEFAULTS_PATH = AGENT_DIR / "config" / "defaults_finetune.json"
 CHECK_CHECKPOINT_PATH = AGENT_DIR / "scripts" / "check_checkpoint.py"
 MAIN_PY_PATH = REPO_ROOT / "main.py"
+FINETUNE_DDP_PATH = REPO_ROOT / "finetune_ddp.py"
 
 # Architecture fields the runner pulls from the ckpt and refuses to let the
 # user override. Mirrors the pretrain runner's ARCH_FLAGS_FROM_CKPT but adds
@@ -189,12 +191,22 @@ 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.
+
+    num_gpus == 1 -> single-process `main.py finetune ...` (default, unchanged).
+    num_gpus  > 1 -> data-parallel `finetune_ddp.py ...` (DDP; the launcher
+    inserts the `finetune` subcommand itself, so it is omitted here). All other
+    flags are identical between the two paths.
+    """
     outputs = manifest["outputs"]
     method = manifest["split_method"]
 
-    argv: list[str] = [sys.executable, "-u", str(MAIN_PY_PATH), "finetune"]
+    if num_gpus > 1:
+        argv: list[str] = [sys.executable, "-u", str(FINETUNE_DDP_PATH)]
+    else:
+        argv = [sys.executable, "-u", str(MAIN_PY_PATH), "finetune"]
 
     # Data + features paths
     if method == "deferred_to_runner":
@@ -273,8 +285,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: finetune_ddp.py 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 +324,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: finetune_ddp.py 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 (main.py finetune, or finetune_ddp.py if num_gpus>1).
     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 +367,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,8 +387,13 @@ def run(args: argparse.Namespace) -> dict[str, Any]:
         return run_manifest
 
     env = os.environ.copy()
-    env["CUDA_VISIBLE_DEVICES"] = str(gpu)
-    # main.py enables strict deterministic algorithms via
+    if distributed:
+        # DDP: finetune_ddp.py spawns one rank per GPU and reads WORLD_SIZE.
+        # 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 / finetune_ddp.py enable strict deterministic algorithms via
     # `torch.use_deterministic_algorithms(True)` (kermt/main.py:23); CuBLAS
     # then requires this env var to be set.
     env.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
@@ -389,8 +419,12 @@ 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 launches "
+                        "finetune_ddp.py with WORLD_SIZE=N; --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 launches `finetune_ddp.py` with
+  `WORLD_SIZE=N`; `--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
 

finetune_ddp.py

@@ -0,0 +1,108 @@
+# 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.
+"""
+Multi-GPU (DistributedDataParallel) launcher for KERMT finetuning.
+
+Mirrors pretrain_ddp.py: one process per GPU via torch.multiprocessing.spawn,
+single-node NCCL rendezvous, data-parallel training with a DistributedSampler.
+The finetune training/eval/checkpointing logic lives in task/train.py; this
+file only bootstraps the process group and hands rank/world_size down through
+cross_validate -> run_training.
+
+Usage (data-parallel finetune across all visible GPUs):
+    python finetune_ddp.py --data_path train.csv --separate_val_path val.csv \
+        --separate_test_path test.csv --checkpoint_path <ckpt> --save_dir <dir> \
+        --dataset_type regression --... (all normal finetune args)
+
+Pin the GPU count explicitly with WORLD_SIZE (defaults to all visible GPUs):
+    WORLD_SIZE=4 python finetune_ddp.py ...
+
+Notes:
+  - `args.batch_size` is the PER-GPU batch size; the effective global batch is
+    batch_size * world_size (same convention as pretrain_ddp.py). The LR
+    schedule is world-size-aware.
+  - Single-process finetune (`python main.py finetune ...`) is unchanged and
+    does not go through this launcher.
+"""
+import os
+import random
+import sys
+
+import numpy as np
+import torch
+import torch.multiprocessing as mp
+from rdkit import RDLogger
+from torch.distributed import destroy_process_group
+
+from kermt.data.torchvocab import MolVocab
+from kermt.util.ddp_utils import configure_nccl_for_topology, ddp_setup
+from kermt.util.parsing import parse_args
+from kermt.util.utils import create_logger
+from task.cross_validate import cross_validate
+
+
+def _setup_determinism(seed: int):
+    """Match the single-process determinism setup in main.py:setup()."""
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.use_deterministic_algorithms(mode=True)
+
+
+def ddp_main(rank: int, world_size: int):
+    # Deterministic algorithms (enabled in _setup_determinism, matching main.py)
+    # require this for CuBLAS on CUDA >= 10.2. Mirrors run_finetune_local.py.
+    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+    ddp_setup(rank, world_size)
+
+    # Suppress RDKit logging in every worker.
+    lg = RDLogger.logger()
+    lg.setLevel(RDLogger.CRITICAL)
+    _ = MolVocab  # ensure vocab class is imported for checkpoint (un)pickling
+
+    # Select the finetune subparser (mirror `main.py finetune ...`). Users invoke
+    # this script with the normal finetune flags but no subcommand token.
+    if len(sys.argv) < 2 or sys.argv[1] != 'finetune':
+        sys.argv.insert(1, 'finetune')
+    args = parse_args()
+
+    _setup_determinism(args.seed)
+
+    # Only rank 0 writes logs; other ranks stay quiet.
+    logger = create_logger(name='train', save_dir=args.save_dir, quiet=False) if rank == 0 else None
+
+    cross_validate(args, logger, rank=rank, world_size=world_size)
+
+    destroy_process_group()
+
+
+if __name__ == '__main__':
+    configure_nccl_for_topology()
+
+    available_gpus = torch.cuda.device_count()
+    if available_gpus == 0:
+        raise RuntimeError('No GPUs found. DDP finetuning requires at least 1 GPU.')
+
+    world_size = int(os.environ.get('WORLD_SIZE', available_gpus))
+    if world_size > available_gpus:
+        raise RuntimeError(
+            f'WORLD_SIZE={world_size} but only {available_gpus} GPU(s) visible. '
+            f'Set WORLD_SIZE<={available_gpus} or unset it to auto-detect.'
+        )
+
+    print(f'Launching DDP finetuning on {world_size}/{available_gpus} GPU(s)')
+    mp.spawn(ddp_main, args=(world_size,), nprocs=world_size)

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)