From 7f9d1853510e1d47d2a4f146d8f5658f6154b4ab Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Tue, 30 Dec 2025 17:34:43 +0000 Subject: [PATCH 01/11] Add SSHLauncher Signed-off-by: Olatunji Ruwase --- apps/grpo/llama3_8b.yaml | 36 ++- apps/grpo/main.py | 4 +- apps/grpo/qwen3_1_7b.yaml | 27 ++- apps/grpo/qwen3_8b.yaml | 39 ++-- src/forge/controller/launchers/__init__.py | 12 + .../controller/launchers/base_launcher.py | 44 ++++ .../controller/launchers/launcher_factory.py | 43 ++++ .../slurm_launcher.py} | 32 +-- .../controller/launchers/ssh_launcher.py | 221 ++++++++++++++++++ src/forge/controller/provisioner.py | 72 +++--- src/forge/types.py | 14 ++ 11 files changed, 448 insertions(+), 96 deletions(-) create mode 100644 src/forge/controller/launchers/__init__.py create mode 100644 src/forge/controller/launchers/base_launcher.py create mode 100644 src/forge/controller/launchers/launcher_factory.py rename src/forge/controller/{launcher.py => launchers/slurm_launcher.py} (74%) create mode 100644 src/forge/controller/launchers/ssh_launcher.py diff --git a/apps/grpo/llama3_8b.yaml b/apps/grpo/llama3_8b.yaml index 6a887ebc3..32062c8a2 100644 --- a/apps/grpo/llama3_8b.yaml +++ b/apps/grpo/llama3_8b.yaml @@ -10,12 +10,19 @@ model: "meta-llama/Meta-Llama-3.1-8B-Instruct" off_by_n: 1 # Off by one by default compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM +# Main loop configuration +rollout_threads: ${services.policy.num_replicas} # Recommended to set equal to policy.num_replicas + +provisioner: + launcher: ssh + ssh_hostfile: "/job/hostfile" + # Observability configuration metric_logging: - wandb: - project: grpo-training - group: grpo_exp_${oc.env:USER} - logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce + # wandb: + # project: grpo-training + # group: grpo_exp_${oc.env:USER} + # logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce console: logging_mode: global_reduce @@ -63,7 +70,7 @@ trainer: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: -1 + data_parallel_shard_degree: 4 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -71,7 +78,7 @@ trainer: disable_loss_parallel: true checkpoint: enable: true - folder: ./checkpoint # The folder to save checkpoints to. + folder: /data-fast/torchforge/${trainer.model.flavor} initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo last_save_in_hf: true @@ -80,6 +87,7 @@ trainer: activation_checkpoint: mode: selective selective_ac_option: op + dcp_path: /modeling-checkpoints/torchforge_dcp_tmp/${trainer.model.flavor} # Replay buffer configuration replay_buffer: @@ -88,7 +96,7 @@ replay_buffer: # This should match the dp_size of TorchTitan # Here it's set explicitly to 2, because we've set # 2 GPUs for the trainer and we're using full FSDP. - dp_size: 2 + dp_size: ${trainer.parallelism.data_parallel_shard_degree} # Must equal trainer DP degree # Reference model configuration ref_model: @@ -104,7 +112,7 @@ ref_model: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: 1 + data_parallel_shard_degree: 4 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -117,14 +125,17 @@ ref_model: services: policy: procs: ${policy.engine_args.tensor_parallel_size} - num_replicas: 1 - with_gpus: true + num_replicas: 2 mesh_name: policy + with_gpus: true + hosts: 1 + ref_model: - procs: 1 + procs: 4 num_replicas: 1 with_gpus: true mesh_name: ref_model + hosts: 1 reward_actor: procs: 1 num_replicas: 1 @@ -137,9 +148,10 @@ actors: with_gpus: false mesh_name: dataset trainer: - procs: 2 + procs: 4 with_gpus: true mesh_name: trainer + hosts: 1 replay_buffer: procs: 1 with_gpus: false diff --git a/apps/grpo/main.py b/apps/grpo/main.py index 224c7b4d5..c08f4f55d 100644 --- a/apps/grpo/main.py +++ b/apps/grpo/main.py @@ -96,8 +96,10 @@ async def main(cfg: DictConfig): # ---- Global setups ---- # provisioner = None if cfg.get("provisioner", None) is not None: + cfg_kwargs = LauncherConfig.prepare_kwargs(cfg) + launcher_config = LauncherConfig(**cfg_kwargs) provisioner = await init_provisioner( - ProvisionerConfig(launcher_config=LauncherConfig(**cfg.provisioner)) + ProvisionerConfig(launcher_config=launcher_config) ) else: provisioner = await init_provisioner() diff --git a/apps/grpo/qwen3_1_7b.yaml b/apps/grpo/qwen3_1_7b.yaml index b0c08a28c..a9bb5b85e 100644 --- a/apps/grpo/qwen3_1_7b.yaml +++ b/apps/grpo/qwen3_1_7b.yaml @@ -11,15 +11,18 @@ off_by_n: 1 # Off by one by default compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM # Main loop configuration -rollout_threads: 1 # Recommended to set equal to policy.num_replicas +rollout_threads: ${services.policy.num_replicas} # Recommended to set equal to policy.num_replicas +provisioner: + launcher: ssh + ssh_hostfile: "/job/hostfile" # Observability configuration metric_logging: - wandb: - project: grpo-training - group: grpo_exp_${oc.env:USER} - logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce + # wandb: + # project: grpo-training + # group: grpo_exp_${oc.env:USER} + # logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce console: logging_mode: global_reduce @@ -67,7 +70,7 @@ trainer: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: 1 + data_parallel_shard_degree: 4 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -75,7 +78,7 @@ trainer: disable_loss_parallel: true checkpoint: enable: true - folder: ./checkpoint # The folder to save checkpoints to. + folder: /data-fast/torchforge/${trainer.model.flavor} initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo last_save_in_hf: true @@ -84,6 +87,7 @@ trainer: activation_checkpoint: mode: selective selective_ac_option: op + dcp_path: /modeling-checkpoints/torchforge_dcp_tmp/${trainer.model.flavor} # Replay buffer configuration replay_buffer: @@ -119,14 +123,16 @@ ref_model: services: policy: procs: ${policy.engine_args.tensor_parallel_size} - num_replicas: 1 + num_replicas: 4 mesh_name: policy with_gpus: true + hosts: 1 ref_model: procs: 1 - num_replicas: 1 + num_replicas: 4 mesh_name: ref_model with_gpus: true + hosts: 1 reward_actor: procs: 1 num_replicas: 1 @@ -139,9 +145,10 @@ actors: with_gpus: false mesh_name: dataset trainer: - procs: 1 + procs: 4 with_gpus: true mesh_name: trainer + hosts: 1 replay_buffer: procs: 1 with_gpus: false diff --git a/apps/grpo/qwen3_8b.yaml b/apps/grpo/qwen3_8b.yaml index 9bec6a541..50807ae58 100644 --- a/apps/grpo/qwen3_8b.yaml +++ b/apps/grpo/qwen3_8b.yaml @@ -10,12 +10,20 @@ model: "Qwen/Qwen3-8B" off_by_n: 1 # Off by one by default compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM +# Main loop configuration +rollout_threads: ${services.policy.num_replicas} # Recommended to set equal to policy.num_replicas + +provisioner: + launcher: ssh + ssh_hostfile: "/job/hostfile" + colocate_ref_and_trainer: true + # Observability configuration metric_logging: - wandb: - project: grpo-training - group: grpo_exp_${oc.env:USER} - logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce + # wandb: + # project: grpo-training + # group: grpo_exp_${oc.env:USER} + # logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce console: logging_mode: global_reduce @@ -33,7 +41,7 @@ policy: model: ${model} tensor_parallel_size: 2 pipeline_parallel_size: 1 - enforce_eager: ${not:${compile}} + enforce_eager: true # ${not:${compile}} sampling_params: # https://docs.vllm.ai/en/v0.10.0/api/vllm/sampling_params.html#vllm.sampling_params.SamplingParams n: ${group_size} max_tokens: ${max_res_tokens} @@ -63,7 +71,7 @@ trainer: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: -1 + data_parallel_shard_degree: 4 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -71,7 +79,7 @@ trainer: disable_loss_parallel: true checkpoint: enable: true - folder: ./checkpoint # The folder to save checkpoints to. + folder: /data-fast/torchforge/${trainer.model.flavor} initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo last_save_in_hf: true @@ -80,6 +88,7 @@ trainer: activation_checkpoint: mode: selective selective_ac_option: op + dcp_path: /modeling-checkpoints/torchforge_dcp_tmp/${trainer.model.flavor} # Replay buffer configuration replay_buffer: @@ -88,7 +97,7 @@ replay_buffer: # This should match the dp_size of TorchTitan # Here it's set explicitly to 2, because we've set # 2 GPUs for the trainer and we're using full FSDP. - dp_size: 2 + dp_size: ${trainer.parallelism.data_parallel_shard_degree} # Must equal trainer DP degree # Reference model configuration ref_model: @@ -104,7 +113,7 @@ ref_model: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: 1 + data_parallel_shard_degree: 4 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -117,14 +126,17 @@ ref_model: services: policy: procs: ${policy.engine_args.tensor_parallel_size} - num_replicas: 1 - with_gpus: true + num_replicas: 2 mesh_name: policy + with_gpus: true + hosts: 1 + ref_model: - procs: 1 + procs: 4 num_replicas: 1 with_gpus: true mesh_name: ref_model + hosts: 1 reward_actor: procs: 1 num_replicas: 1 @@ -137,9 +149,10 @@ actors: with_gpus: false mesh_name: dataset trainer: - procs: 2 + procs: 4 with_gpus: true mesh_name: trainer + hosts: 1 replay_buffer: procs: 1 with_gpus: false diff --git a/src/forge/controller/launchers/__init__.py b/src/forge/controller/launchers/__init__.py new file mode 100644 index 000000000..514152651 --- /dev/null +++ b/src/forge/controller/launchers/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from .base_launcher import BaseLauncher +from .launcher_factory import get_launcher + +__all__ = [ + "BaseLauncher", + "get_launcher", +] diff --git a/src/forge/controller/launchers/base_launcher.py b/src/forge/controller/launchers/base_launcher.py new file mode 100644 index 000000000..1e853a98a --- /dev/null +++ b/src/forge/controller/launchers/base_launcher.py @@ -0,0 +1,44 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Base Launcher class""" + +from typing import Any +import monarch +from monarch.actor import ProcMesh, HostMesh + + +class BaseLauncher: + async def initialize(self) -> None: + pass + + async def get_allocator(self, name: str, num_hosts: int) -> tuple[Any, Any, str]: + pass + + async def get_host_mesh(self, name: str, num_hosts: int) -> tuple[HostMesh, str]: + alloc, alloc_constraints, server_name = await self.get_allocator( + name, num_hosts + ) + + # We are asking Monarch to allocate a single process on + # every host, reflected in the Extent we provide below. + + # Technically, this is ["hosts", "procs"] but to reduce + # confusion on its relationship with procs elsewhere, + # we call it "no_dim". + + # TODO - remove this once Monarch supports HostMesh without it. + host_mesh = HostMesh.allocate_nonblocking( + name=name, + extent=Extent(["hosts", "no_dim"], [num_hosts, 1]), + allocator=alloc, + alloc_constraints=alloc_constraints, + ) + return host_mesh, server_name + + async def remote_setup(self, procs: ProcMesh) -> None: + pass + diff --git a/src/forge/controller/launchers/launcher_factory.py b/src/forge/controller/launchers/launcher_factory.py new file mode 100644 index 000000000..3ecdebf2e --- /dev/null +++ b/src/forge/controller/launchers/launcher_factory.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Launcher Factory""" + +import os +from typing import Any + +from forge.types import Launcher, LauncherConfig +from .base_launcher import BaseLauncher +from .slurm_launcher import Slurmlauncher +from .ssh_launcher import SSHLauncher + +JOB_NAME_KEY = "job_name" +LAUNCHER_KEY = "launcher" + +def get_launcher(cfg: LauncherConfig | None = None) -> BaseLauncher | None: + if not cfg: + return None + if cfg.launcher == Launcher.SLURM: + return Slurmlauncher(cfg) + elif cfg.launcher == Launcher.SSH: + success, error_msg = SSHLauncher.validate_configuration(cfg) + if not success: + raise ValueError(error_msg) + # if cfg.ssh_hostfile is None or not os.path.isfile(cfg.ssh_hostfile): + # raise ValueError( + # f"SSH Launcher cannot work because of invalid hostfile path: {cfg.ssh_hostfile}." + # ) + return SSHLauncher(cfg) + elif cfg.launcher == Launcher.MAST: + try: + from forge.fb.mast_launcher import MastLauncher + + return MastLauncher(cfg, detached=False) + except ImportError as err: + raise ValueError("MAST is not available, cannot launch MAST jobs.") from err + + else: + raise ValueError(f"Unsupported config provided, got {cfg}") \ No newline at end of file diff --git a/src/forge/controller/launcher.py b/src/forge/controller/launchers/slurm_launcher.py similarity index 74% rename from src/forge/controller/launcher.py rename to src/forge/controller/launchers/slurm_launcher.py index 7ce38d71f..f5621b054 100644 --- a/src/forge/controller/launcher.py +++ b/src/forge/controller/launchers/slurm_launcher.py @@ -4,27 +4,26 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Launcher specific logic (i.e. SLURM, k8s when supported, etc.)""" +"""Slurm Launcher""" +import os +import sys import tempfile from typing import Any import monarch -from forge.controller.base import BaseLauncher -from forge.types import Launcher, LauncherConfig + +from forge.types import LauncherConfig +from forge.controller.launchers.base_launcher import BaseLauncher + from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport from monarch._rust_bindings.monarch_hyperactor.config import configure from monarch._src.actor.allocator import RemoteAllocator, TorchXRemoteAllocInitializer from monarch.actor import ProcMesh from monarch.tools import commands -from monarch.tools.components import hyperactor from monarch.tools.config import Config -JOB_NAME_KEY = "job_name" -LAUNCHER_KEY = "launcher" - - class Slurmlauncher(BaseLauncher): def __init__( self, @@ -73,20 +72,3 @@ async def get_allocator(self, name: str, num_hosts: int) -> tuple[Any, Any, str] async def remote_setup(self, procs: ProcMesh) -> None: return - - -def get_launcher(cfg: LauncherConfig | None = None) -> BaseLauncher | None: - if not cfg: - return None - if cfg.launcher == Launcher.SLURM: - return Slurmlauncher(cfg) - elif cfg.launcher == Launcher.MAST: - try: - from forge.fb.mast_launcher import MastLauncher - - return MastLauncher(cfg, detached=False) - except ImportError as err: - raise ValueError("MAST is not available, cannot launch MAST jobs.") from err - - else: - raise ValueError(f"Unsupported config provided, got {cfg}") diff --git a/src/forge/controller/launchers/ssh_launcher.py b/src/forge/controller/launchers/ssh_launcher.py new file mode 100644 index 000000000..e5a0c24fd --- /dev/null +++ b/src/forge/controller/launchers/ssh_launcher.py @@ -0,0 +1,221 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""SSHLauncher """ + +import os +import sys +import subprocess +from typing import Any +import collections +import re + +from dataclasses import dataclass, field +from typing import List + +from forge.types import LauncherConfig + +from .base_launcher import BaseLauncher + +import monarch +from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport +from monarch._rust_bindings.monarch_hyperactor.config import configure +from monarch.actor import Actor, endpoint, ProcMesh, HostMesh +from monarch.tools import commands +from monarch.tools.commands import create, info +from monarch._src.job.job import SSHJob + +def fetch_hostfile(hostfile_path): + # e.g., worker-0 slots=16 + with open(hostfile_path, 'r') as fd: + hostfile_text = fd.readlines() + + return parse_hostfile(hostfile_text) + + +def parse_hostfile(hostfile_lines): + # Regex matches one or more non-whitespace characters (\S+) at the start of + # the line, followed by one or more whitespace characters (\s+), followed + # by the string "slots=", followed by one or more digits (\d+). + pattern = r'^(\S+)\s+slots=(\d+)' + + resource_pool = collections.OrderedDict() + + for line in hostfile_lines: + line = line.strip() + match = re.search(pattern, line) + if line.startswith("#") or line == "": + # hostfile comment or empty line, ignore + continue + elif match: + host = match.group(1) + num_slots = int(match.group(2)) + if host in resource_pool: + raise ValueError(f"Hostfile contains multiple entries for {host}, unable to proceed with launching") + resource_pool[host] = num_slots + else: + raise ValueError(f"Hostfile contains a bad entry: {line}, unable to proceed with launching") + + if len(resource_pool) == 0: + raise ValueError("Hostfile is empty or not formatted correctly, unable to proceed with launching.") + + return resource_pool + +POLICY_KEY = "policy" +REFERENCE_KEY = "ref_model" +TRAINER_KEY = "trainer" +ProcMesh_KEYS = [POLICY_KEY, REFERENCE_KEY, TRAINER_KEY] + + +@dataclass +class SSHForgeProc: + is_actor: bool = True + num_hosts: int = 0 + +@dataclass +class SSHActor: + name: str | None = None + host_mesh: HostMesh | None = None + + def get_host_mesh(self): + return self.host_mesh + +@dataclass +class SSHService: + name: str | None = None + host_mesh_list: List[HostMesh] = field(default_factory=list) + current_idx: int = 0 + + def get_host_mesh(self): + host_mesh = self.host_mesh_list[self.current_idx] + self.current_idx = (self.current_idx + 1) % len(self.host_mesh_list) + return host_mesh + +class SSHLauncher(BaseLauncher): + def __init__(self, cfg: LauncherConfig): + self.job = None + self.state = None + self.ssh_hostfile = cfg.ssh_hostfile + self.host_pool = fetch_hostfile(self.ssh_hostfile) + self.host_ips = list(self.host_pool.keys()) + + self.forge_proc_map = self._create_forge_proc_map(cfg) + n_policy_slice_end = self.forge_proc_map[POLICY_KEY].num_hosts + n_trainer_slice_end = n_policy_slice_end + self.forge_proc_map[TRAINER_KEY].num_hosts + n_ref_model_slice_start = n_policy_slice_end if cfg.colocate_ref_and_trainer else n_trainer_slice_end + n_ref_model_slice_end = n_ref_model_slice_start + self.forge_proc_map[REFERENCE_KEY].num_hosts + + self.host_ip_map = { + POLICY_KEY: self.host_ips[:n_policy_slice_end], + TRAINER_KEY: self.host_ips[n_policy_slice_end: n_trainer_slice_end], + REFERENCE_KEY: self.host_ips[n_ref_model_slice_start: n_ref_model_slice_end] + } + + self.host_mesh_map = dict() + + + def _create_forge_proc_map(self, cfg: LauncherConfig) -> dict: + proc_map = {} + for key in ProcMesh_KEYS: + if key in cfg.services: + proc_map[key] = SSHForgeProc(is_actor=False, num_hosts=cfg.services[key].hosts) + else: + proc_map[key] = SSHForgeProc(is_actor=True, num_hosts=cfg.actors[key].hosts) + + return proc_map + + async def _setup_job(self): + job = SSHJob( + python_exe=os.getenv("PYTHON_EXE", sys.executable), + ssh_args=[ + "-o", "StrictHostKeyChecking=no", + "-o", "BatchMode=yes", + ], + monarch_port=int(os.getenv("MONARCH_PORT", "22222")), # make sure this is open + ) + + for proc_key in ProcMesh_KEYS: + job.add_mesh(f"{proc_key}_mesh", self.host_ip_map[proc_key]) + + return job + + + async def get_host_mesh(self, name: str, num_hosts: int) -> tuple[HostMesh, str]: + print(f"ssh_launcher_get_host_mesh: {name=} {num_hosts=}") + for key in self.host_mesh_map.keys(): + if key in name: + host_mesh = self.host_mesh_map[key].get_host_mesh() + return host_mesh, None + + return None, None + + + def _create_actor_or_service(self, mesh_name, job_host_mesh): + if self.forge_proc_map[mesh_name].is_actor: + return SSHActor(name=mesh_name, host_mesh=job_host_mesh) + + return SSHService( + name = mesh_name, + host_mesh_list = [job_host_mesh.slice(hosts=slice(i,i+1)) for i in range(len(self.host_ip_map[mesh_name]))] + ) + + async def initialize(self) -> None: + # HostMesh currently requires explicit configuration + # of the underlying transport from client to mesh. + # This can be removed in the future once this has been removed. + configure(default_transport=ChannelTransport.TcpWithHostname) + + self.job = await self._setup_job() + self.state = self.job.state() + + self.host_mesh_map[TRAINER_KEY] = self._create_actor_or_service(mesh_name=TRAINER_KEY, job_host_mesh=self.state.trainer_mesh) + self.host_mesh_map[POLICY_KEY] = self._create_actor_or_service(mesh_name=POLICY_KEY, job_host_mesh=self.state.policy_mesh) + self.host_mesh_map[REFERENCE_KEY] = self._create_actor_or_service(mesh_name=REFERENCE_KEY, job_host_mesh=self.state.ref_model_mesh) + + print(f'ssh_launcher[init]: {self.ssh_hostfile=} {self.host_ips=}') + print(f'ssh_launcher[init]: {self.host_ip_map=}') + + + @classmethod + def validate_configuration(cls, cfg: LauncherConfig) -> tuple[bool, str]: + + error_msg_prefix = "SSH Launcher configuration error:" + # Check that a valid hostfile is specified. + if cfg.ssh_hostfile is None or not os.path.isfile(cfg.ssh_hostfile): + error_msg = f"{error_msg_prefix} Invalid hostfile path {cfg.ssh_hostfile}." + return False, error_msg + + # Check that all services/actors which use gpus also specify host count + print(f"{cfg=}") + print(f"{cfg.services=}") + print(f"{cfg.actors=}") + + gpu_host_count = {name: service.hosts for name, service in cfg.services.items() if service.with_gpus} + gpu_host_count.update({name: actor.hosts for name, actor in cfg.actors.items() if actor.with_gpus}) + missing_hosts = [name for name, count in gpu_host_count.items() if count is None or count < 1] + if len(missing_hosts) > 0: + error_msg = f"{error_msg_prefix} missing host counts for services/actors that need gpus {missing_hosts}" + return False, error_msg + + # Check that sufficient hosts available + available_host_count = len(fetch_hostfile(cfg.ssh_hostfile)) + required_host_count = sum(list(gpu_host_count.values())) + + if cfg.colocate_ref_and_trainer and TRAINER_KEY in gpu_host_count and REFERENCE_KEY in gpu_host_count: + excess_gpu_count = min(gpu_host_count[TRAINER_KEY], gpu_host_count[REFERENCE_KEY]) + required_host_count -= excess_gpu_count + + if required_host_count > available_host_count: + error_msg = f"{error_msg_prefix} {required_host_count=} exceeds {available_host_count=}" + return False, error_msg + + print(f"{required_host_count=} {available_host_count=} {gpu_host_count=}") + + return True, None + + + + diff --git a/src/forge/controller/provisioner.py b/src/forge/controller/provisioner.py index b7b703841..1ba234c29 100644 --- a/src/forge/controller/provisioner.py +++ b/src/forge/controller/provisioner.py @@ -15,7 +15,7 @@ import torch -from forge.controller.launcher import BaseLauncher, get_launcher +from forge.controller.launchers import BaseLauncher, get_launcher from forge.env import all_env_vars, FORGE_DISABLE_METRICS from forge.types import ProcessConfig, ProvisionerConfig @@ -179,10 +179,10 @@ def get_available_gpus(self) -> list[str]: """Returns a list of available GPU devices.""" return [str(gpu) for gpu in self.available_gpus] - def get_gpus(self, num_gpus: int) -> list[str]: + def get_gpus(self, num_gpus: int, mesh_name: str) -> list[str]: """Assigns GPU devices.""" if num_gpus > len(self.available_gpus): - raise RuntimeError("Not enough GPUs available") + raise RuntimeError(f"Not enough GPUs available: {mesh_name=} request_num={num_gpus}, available_num={len(self.available_gpus)} {self.available_gpus=}") gpus = list(self.available_gpus)[:num_gpus] self.available_gpus -= set(gpus) return [str(gpu) for gpu in gpus] @@ -260,24 +260,7 @@ async def create_host_mesh(self, name: str, num_hosts: int) -> HostMesh: "but no launcher was specified." ) logger.debug(f"Creating remote server for alloc {name}") - alloc, alloc_constraints, server_name = await self.launcher.get_allocator( - name, num_hosts - ) - - # We are asking Monarch to allocate a single process on - # every host, reflected in the Extent we provide below. - - # Technically, this is ["hosts", "procs"] but to reduce - # confusion on its relationship with procs elsewhere, - # we call it "no_dim". - - # TODO - remove this once Monarch supports HostMesh without it. - host_mesh = HostMesh.allocate_nonblocking( - name=name, - extent=Extent(["hosts", "no_dim"], [num_hosts, 1]), - allocator=alloc, - alloc_constraints=alloc_constraints, - ) + host_mesh, server_name = await self.launcher.get_host_mesh() return host_mesh, server_name def get_host_mesh(self, name: str) -> HostMesh: @@ -289,6 +272,25 @@ def get_host_mesh(self, name: str) -> HostMesh: """ return self._host_mesh_map[name] + + async def add_host_mesh(self, host_mesh: HostMesh, mesh_name: str) -> None: + if host_mesh is None: + raise RuntimeError( + "Cannot add host_mesh that is None, " + ) + + if hasattr(host_mesh, '_host_id') and host_mesh._host_id in self._host_gpu_map: + logger.debug(f'add_host_mesh: duplicate host_mesh {host_mesh._host_id=}') + return + + host_id = uuid.uuid1() + # Get the GPU count from the remote host + remote_gpu_count = await get_host_gpus(host_mesh) + gpu_manager = GpuManager(max_device_count=remote_gpu_count) + self._host_gpu_map[host_id] = gpu_manager + host_mesh._host_id = host_id + + async def get_proc_mesh( self, num_procs: int, @@ -324,8 +326,8 @@ async def get_proc_mesh( if env_vars is None: env_vars = {} - is_remote = num_hosts is not None and num_hosts > 0 - + is_remote = num_hosts is not None and num_hosts > 0 and self.launcher is not None + host_mesh_extent = host_mesh.extent if host_mesh else None async with self._lock: server_name = None if is_remote: @@ -333,29 +335,27 @@ async def get_proc_mesh( created_hosts = len(self._server_names) mesh_name = f"alloc_{created_hosts}" if host_mesh is None: - host_mesh, server_name = await self.create_host_mesh( + host_mesh, server_name = await self.launcher.get_host_mesh( name=mesh_name, num_hosts=num_hosts, ) - host_id = uuid.uuid1() - # Get the GPU count from the remote host - remote_gpu_count = await get_host_gpus(host_mesh) - gpu_manager = GpuManager(max_device_count=remote_gpu_count) - self._host_gpu_map[host_id] = gpu_manager - host_mesh._host_id = host_id - else: - host_id = host_mesh._host_id - gpu_manager = self._host_gpu_map[host_id] + + assert host_mesh, f"host_mesh cannot be None at this point." + await self.add_host_mesh(host_mesh, mesh_name) + + host_id = host_mesh._host_id + gpu_manager = self._host_gpu_map[host_id] else: # fallback to local host_mesh = this_host() gpu_manager = self._host_gpu_map[self._this_host_id] host_mesh._host_id = self._this_host_id + per_host_spec = {"procs": num_procs} if with_gpus: if not addr or not port: addr, port = await get_remote_info(host_mesh) - gpu_ids = gpu_manager.get_gpus(num_procs) + gpu_ids = gpu_manager.get_gpus(num_procs, mesh_name) env_vars["MASTER_ADDR"] = addr env_vars["MASTER_PORT"] = port @@ -369,9 +369,11 @@ async def get_proc_mesh( for env_var in all_env_vars(): env_vars[env_var.name] = str(env_var.get_value()) + per_host_spec = {"gpus": len(gpu_ids)} + # Spawn procs without bootstrap to avoid SetupActor mesh failures procs = host_mesh.spawn_procs( - per_host={"procs": num_procs}, + per_host=per_host_spec, name=mesh_name, ) diff --git a/src/forge/types.py b/src/forge/types.py index 8ba17bd01..3d91bd502 100644 --- a/src/forge/types.py +++ b/src/forge/types.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from enum import Enum from typing import Any, TypedDict, Union +from omegaconf import DictConfig class Message(TypedDict): @@ -38,6 +39,7 @@ class Observation: class Launcher(Enum): MAST = "mast" SLURM = "slurm" + SSH = "ssh" @dataclass @@ -116,11 +118,23 @@ class LauncherConfig: gpu: int = 8 # GPUs per node (required for SLURM, can get with sinfo) account: str = "" qos: str = "" + ssh_hostfile: str = None + colocate_ref_and_trainer: bool = True def __post_init__(self): if isinstance(self.launcher, str): self.launcher = Launcher(self.launcher) + self.services = {name: ServiceConfig(**value) for name, value in self.services.items()} + self.actors = {name: ProcessConfig(**value) for name, value in self.actors.items()} + + @classmethod + def prepare_kwargs(cls, cfg: DictConfig) -> dict: + kwargs = {} + kwargs.update(**cfg.provisioner) + kwargs["services"] = cfg.services + kwargs["actors"] = cfg.actors + return kwargs @dataclass class ProvisionerConfig: From 7786e217bdf64dc9474c516d7acdc86728607785 Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Tue, 30 Dec 2025 17:46:12 +0000 Subject: [PATCH 02/11] Formatting Signed-off-by: Olatunji Ruwase --- .../controller/launchers/base_launcher.py | 7 +- .../controller/launchers/launcher_factory.py | 6 +- .../controller/launchers/slurm_launcher.py | 5 +- .../controller/launchers/ssh_launcher.py | 158 +++++++++++------- src/forge/controller/provisioner.py | 32 ++-- src/forge/types.py | 12 +- 6 files changed, 130 insertions(+), 90 deletions(-) diff --git a/src/forge/controller/launchers/base_launcher.py b/src/forge/controller/launchers/base_launcher.py index 1e853a98a..be7202f0c 100644 --- a/src/forge/controller/launchers/base_launcher.py +++ b/src/forge/controller/launchers/base_launcher.py @@ -7,8 +7,8 @@ """Base Launcher class""" from typing import Any -import monarch -from monarch.actor import ProcMesh, HostMesh + +from monarch.actor import HostMesh, ProcMesh class BaseLauncher: @@ -37,8 +37,7 @@ async def get_host_mesh(self, name: str, num_hosts: int) -> tuple[HostMesh, str] allocator=alloc, alloc_constraints=alloc_constraints, ) - return host_mesh, server_name + return host_mesh, server_name async def remote_setup(self, procs: ProcMesh) -> None: pass - diff --git a/src/forge/controller/launchers/launcher_factory.py b/src/forge/controller/launchers/launcher_factory.py index 3ecdebf2e..47711c875 100644 --- a/src/forge/controller/launchers/launcher_factory.py +++ b/src/forge/controller/launchers/launcher_factory.py @@ -6,10 +6,9 @@ """Launcher Factory""" -import os -from typing import Any from forge.types import Launcher, LauncherConfig + from .base_launcher import BaseLauncher from .slurm_launcher import Slurmlauncher from .ssh_launcher import SSHLauncher @@ -17,6 +16,7 @@ JOB_NAME_KEY = "job_name" LAUNCHER_KEY = "launcher" + def get_launcher(cfg: LauncherConfig | None = None) -> BaseLauncher | None: if not cfg: return None @@ -40,4 +40,4 @@ def get_launcher(cfg: LauncherConfig | None = None) -> BaseLauncher | None: raise ValueError("MAST is not available, cannot launch MAST jobs.") from err else: - raise ValueError(f"Unsupported config provided, got {cfg}") \ No newline at end of file + raise ValueError(f"Unsupported config provided, got {cfg}") diff --git a/src/forge/controller/launchers/slurm_launcher.py b/src/forge/controller/launchers/slurm_launcher.py index f5621b054..2c41075f4 100644 --- a/src/forge/controller/launchers/slurm_launcher.py +++ b/src/forge/controller/launchers/slurm_launcher.py @@ -6,15 +6,14 @@ """Slurm Launcher""" -import os -import sys + import tempfile from typing import Any import monarch +from forge.controller.launchers.base_launcher import BaseLauncher from forge.types import LauncherConfig -from forge.controller.launchers.base_launcher import BaseLauncher from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport from monarch._rust_bindings.monarch_hyperactor.config import configure diff --git a/src/forge/controller/launchers/ssh_launcher.py b/src/forge/controller/launchers/ssh_launcher.py index e5a0c24fd..f4ed7c046 100644 --- a/src/forge/controller/launchers/ssh_launcher.py +++ b/src/forge/controller/launchers/ssh_launcher.py @@ -6,31 +6,26 @@ """SSHLauncher """ +import collections import os -import sys -import subprocess -from typing import Any -import collections -import re +import re +import sys from dataclasses import dataclass, field from typing import List from forge.types import LauncherConfig - -from .base_launcher import BaseLauncher - -import monarch from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport from monarch._rust_bindings.monarch_hyperactor.config import configure -from monarch.actor import Actor, endpoint, ProcMesh, HostMesh -from monarch.tools import commands -from monarch.tools.commands import create, info from monarch._src.job.job import SSHJob +from monarch.actor import HostMesh + +from .base_launcher import BaseLauncher + def fetch_hostfile(hostfile_path): # e.g., worker-0 slots=16 - with open(hostfile_path, 'r') as fd: + with open(hostfile_path, "r") as fd: hostfile_text = fd.readlines() return parse_hostfile(hostfile_text) @@ -40,7 +35,7 @@ def parse_hostfile(hostfile_lines): # Regex matches one or more non-whitespace characters (\S+) at the start of # the line, followed by one or more whitespace characters (\s+), followed # by the string "slots=", followed by one or more digits (\d+). - pattern = r'^(\S+)\s+slots=(\d+)' + pattern = r"^(\S+)\s+slots=(\d+)" resource_pool = collections.OrderedDict() @@ -54,16 +49,23 @@ def parse_hostfile(hostfile_lines): host = match.group(1) num_slots = int(match.group(2)) if host in resource_pool: - raise ValueError(f"Hostfile contains multiple entries for {host}, unable to proceed with launching") + raise ValueError( + f"Hostfile contains multiple entries for {host}, unable to proceed with launching" + ) resource_pool[host] = num_slots else: - raise ValueError(f"Hostfile contains a bad entry: {line}, unable to proceed with launching") + raise ValueError( + f"Hostfile contains a bad entry: {line}, unable to proceed with launching" + ) if len(resource_pool) == 0: - raise ValueError("Hostfile is empty or not formatted correctly, unable to proceed with launching.") + raise ValueError( + "Hostfile is empty or not formatted correctly, unable to proceed with launching." + ) return resource_pool + POLICY_KEY = "policy" REFERENCE_KEY = "ref_model" TRAINER_KEY = "trainer" @@ -75,55 +77,67 @@ class SSHForgeProc: is_actor: bool = True num_hosts: int = 0 + @dataclass class SSHActor: name: str | None = None - host_mesh: HostMesh | None = None + host_mesh: HostMesh | None = None def get_host_mesh(self): return self.host_mesh -@dataclass + +@dataclass class SSHService: - name: str | None = None + name: str | None = None host_mesh_list: List[HostMesh] = field(default_factory=list) - current_idx: int = 0 + current_idx: int = 0 def get_host_mesh(self): host_mesh = self.host_mesh_list[self.current_idx] self.current_idx = (self.current_idx + 1) % len(self.host_mesh_list) return host_mesh + class SSHLauncher(BaseLauncher): def __init__(self, cfg: LauncherConfig): - self.job = None - self.state = None + self.job = None + self.state = None self.ssh_hostfile = cfg.ssh_hostfile self.host_pool = fetch_hostfile(self.ssh_hostfile) self.host_ips = list(self.host_pool.keys()) self.forge_proc_map = self._create_forge_proc_map(cfg) n_policy_slice_end = self.forge_proc_map[POLICY_KEY].num_hosts - n_trainer_slice_end = n_policy_slice_end + self.forge_proc_map[TRAINER_KEY].num_hosts - n_ref_model_slice_start = n_policy_slice_end if cfg.colocate_ref_and_trainer else n_trainer_slice_end - n_ref_model_slice_end = n_ref_model_slice_start + self.forge_proc_map[REFERENCE_KEY].num_hosts - + n_trainer_slice_end = ( + n_policy_slice_end + self.forge_proc_map[TRAINER_KEY].num_hosts + ) + n_ref_model_slice_start = ( + n_policy_slice_end if cfg.colocate_ref_and_trainer else n_trainer_slice_end + ) + n_ref_model_slice_end = ( + n_ref_model_slice_start + self.forge_proc_map[REFERENCE_KEY].num_hosts + ) + self.host_ip_map = { POLICY_KEY: self.host_ips[:n_policy_slice_end], - TRAINER_KEY: self.host_ips[n_policy_slice_end: n_trainer_slice_end], - REFERENCE_KEY: self.host_ips[n_ref_model_slice_start: n_ref_model_slice_end] + TRAINER_KEY: self.host_ips[n_policy_slice_end:n_trainer_slice_end], + REFERENCE_KEY: self.host_ips[n_ref_model_slice_start:n_ref_model_slice_end], } self.host_mesh_map = dict() - def _create_forge_proc_map(self, cfg: LauncherConfig) -> dict: proc_map = {} for key in ProcMesh_KEYS: if key in cfg.services: - proc_map[key] = SSHForgeProc(is_actor=False, num_hosts=cfg.services[key].hosts) + proc_map[key] = SSHForgeProc( + is_actor=False, num_hosts=cfg.services[key].hosts + ) else: - proc_map[key] = SSHForgeProc(is_actor=True, num_hosts=cfg.actors[key].hosts) + proc_map[key] = SSHForgeProc( + is_actor=True, num_hosts=cfg.actors[key].hosts + ) return proc_map @@ -131,35 +145,40 @@ async def _setup_job(self): job = SSHJob( python_exe=os.getenv("PYTHON_EXE", sys.executable), ssh_args=[ - "-o", "StrictHostKeyChecking=no", - "-o", "BatchMode=yes", + "-o", + "StrictHostKeyChecking=no", + "-o", + "BatchMode=yes", ], - monarch_port=int(os.getenv("MONARCH_PORT", "22222")), # make sure this is open + monarch_port=int( + os.getenv("MONARCH_PORT", "22222") + ), # make sure this is open ) for proc_key in ProcMesh_KEYS: job.add_mesh(f"{proc_key}_mesh", self.host_ip_map[proc_key]) - return job - + return job async def get_host_mesh(self, name: str, num_hosts: int) -> tuple[HostMesh, str]: print(f"ssh_launcher_get_host_mesh: {name=} {num_hosts=}") for key in self.host_mesh_map.keys(): if key in name: host_mesh = self.host_mesh_map[key].get_host_mesh() - return host_mesh, None - - return None, None + return host_mesh, None + return None, None def _create_actor_or_service(self, mesh_name, job_host_mesh): if self.forge_proc_map[mesh_name].is_actor: return SSHActor(name=mesh_name, host_mesh=job_host_mesh) return SSHService( - name = mesh_name, - host_mesh_list = [job_host_mesh.slice(hosts=slice(i,i+1)) for i in range(len(self.host_ip_map[mesh_name]))] + name=mesh_name, + host_mesh_list=[ + job_host_mesh.slice(hosts=slice(i, i + 1)) + for i in range(len(self.host_ip_map[mesh_name])) + ], ) async def initialize(self) -> None: @@ -171,21 +190,26 @@ async def initialize(self) -> None: self.job = await self._setup_job() self.state = self.job.state() - self.host_mesh_map[TRAINER_KEY] = self._create_actor_or_service(mesh_name=TRAINER_KEY, job_host_mesh=self.state.trainer_mesh) - self.host_mesh_map[POLICY_KEY] = self._create_actor_or_service(mesh_name=POLICY_KEY, job_host_mesh=self.state.policy_mesh) - self.host_mesh_map[REFERENCE_KEY] = self._create_actor_or_service(mesh_name=REFERENCE_KEY, job_host_mesh=self.state.ref_model_mesh) + self.host_mesh_map[TRAINER_KEY] = self._create_actor_or_service( + mesh_name=TRAINER_KEY, job_host_mesh=self.state.trainer_mesh + ) + self.host_mesh_map[POLICY_KEY] = self._create_actor_or_service( + mesh_name=POLICY_KEY, job_host_mesh=self.state.policy_mesh + ) + self.host_mesh_map[REFERENCE_KEY] = self._create_actor_or_service( + mesh_name=REFERENCE_KEY, job_host_mesh=self.state.ref_model_mesh + ) - print(f'ssh_launcher[init]: {self.ssh_hostfile=} {self.host_ips=}') - print(f'ssh_launcher[init]: {self.host_ip_map=}') - + print(f"ssh_launcher[init]: {self.ssh_hostfile=} {self.host_ips=}") + print(f"ssh_launcher[init]: {self.host_ip_map=}") @classmethod def validate_configuration(cls, cfg: LauncherConfig) -> tuple[bool, str]: error_msg_prefix = "SSH Launcher configuration error:" - # Check that a valid hostfile is specified. + # Check that a valid hostfile is specified. if cfg.ssh_hostfile is None or not os.path.isfile(cfg.ssh_hostfile): - error_msg = f"{error_msg_prefix} Invalid hostfile path {cfg.ssh_hostfile}." + error_msg = f"{error_msg_prefix} Invalid hostfile path {cfg.ssh_hostfile}." return False, error_msg # Check that all services/actors which use gpus also specify host count @@ -193,9 +217,17 @@ def validate_configuration(cls, cfg: LauncherConfig) -> tuple[bool, str]: print(f"{cfg.services=}") print(f"{cfg.actors=}") - gpu_host_count = {name: service.hosts for name, service in cfg.services.items() if service.with_gpus} - gpu_host_count.update({name: actor.hosts for name, actor in cfg.actors.items() if actor.with_gpus}) - missing_hosts = [name for name, count in gpu_host_count.items() if count is None or count < 1] + gpu_host_count = { + name: service.hosts + for name, service in cfg.services.items() + if service.with_gpus + } + gpu_host_count.update( + {name: actor.hosts for name, actor in cfg.actors.items() if actor.with_gpus} + ) + missing_hosts = [ + name for name, count in gpu_host_count.items() if count is None or count < 1 + ] if len(missing_hosts) > 0: error_msg = f"{error_msg_prefix} missing host counts for services/actors that need gpus {missing_hosts}" return False, error_msg @@ -203,19 +235,21 @@ def validate_configuration(cls, cfg: LauncherConfig) -> tuple[bool, str]: # Check that sufficient hosts available available_host_count = len(fetch_hostfile(cfg.ssh_hostfile)) required_host_count = sum(list(gpu_host_count.values())) - - if cfg.colocate_ref_and_trainer and TRAINER_KEY in gpu_host_count and REFERENCE_KEY in gpu_host_count: - excess_gpu_count = min(gpu_host_count[TRAINER_KEY], gpu_host_count[REFERENCE_KEY]) + + if ( + cfg.colocate_ref_and_trainer + and TRAINER_KEY in gpu_host_count + and REFERENCE_KEY in gpu_host_count + ): + excess_gpu_count = min( + gpu_host_count[TRAINER_KEY], gpu_host_count[REFERENCE_KEY] + ) required_host_count -= excess_gpu_count - + if required_host_count > available_host_count: error_msg = f"{error_msg_prefix} {required_host_count=} exceeds {available_host_count=}" return False, error_msg print(f"{required_host_count=} {available_host_count=} {gpu_host_count=}") - return True, None - - - - + return True, None diff --git a/src/forge/controller/provisioner.py b/src/forge/controller/provisioner.py index 1ba234c29..db1bace31 100644 --- a/src/forge/controller/provisioner.py +++ b/src/forge/controller/provisioner.py @@ -20,7 +20,6 @@ from forge.types import ProcessConfig, ProvisionerConfig from monarch._src.actor.actor_mesh import ActorMesh -from monarch._src.actor.shape import Extent from monarch.actor import ( Actor, @@ -182,7 +181,10 @@ def get_available_gpus(self) -> list[str]: def get_gpus(self, num_gpus: int, mesh_name: str) -> list[str]: """Assigns GPU devices.""" if num_gpus > len(self.available_gpus): - raise RuntimeError(f"Not enough GPUs available: {mesh_name=} request_num={num_gpus}, available_num={len(self.available_gpus)} {self.available_gpus=}") + raise RuntimeError( + f"Not enough GPUs available: {mesh_name=} request_num={num_gpus}, " + f"available_num={len(self.available_gpus)} {self.available_gpus=}" + ) gpus = list(self.available_gpus)[:num_gpus] self.available_gpus -= set(gpus) return [str(gpu) for gpu in gpus] @@ -260,7 +262,7 @@ async def create_host_mesh(self, name: str, num_hosts: int) -> HostMesh: "but no launcher was specified." ) logger.debug(f"Creating remote server for alloc {name}") - host_mesh, server_name = await self.launcher.get_host_mesh() + host_mesh, server_name = await self.launcher.get_host_mesh() return host_mesh, server_name def get_host_mesh(self, name: str) -> HostMesh: @@ -272,16 +274,13 @@ def get_host_mesh(self, name: str) -> HostMesh: """ return self._host_mesh_map[name] - async def add_host_mesh(self, host_mesh: HostMesh, mesh_name: str) -> None: if host_mesh is None: - raise RuntimeError( - "Cannot add host_mesh that is None, " - ) + raise RuntimeError("Cannot add host_mesh that is None, ") - if hasattr(host_mesh, '_host_id') and host_mesh._host_id in self._host_gpu_map: - logger.debug(f'add_host_mesh: duplicate host_mesh {host_mesh._host_id=}') - return + if hasattr(host_mesh, "_host_id") and host_mesh._host_id in self._host_gpu_map: + logger.debug(f"add_host_mesh: duplicate host_mesh {host_mesh._host_id=}") + return host_id = uuid.uuid1() # Get the GPU count from the remote host @@ -290,7 +289,6 @@ async def add_host_mesh(self, host_mesh: HostMesh, mesh_name: str) -> None: self._host_gpu_map[host_id] = gpu_manager host_mesh._host_id = host_id - async def get_proc_mesh( self, num_procs: int, @@ -326,8 +324,10 @@ async def get_proc_mesh( if env_vars is None: env_vars = {} - is_remote = num_hosts is not None and num_hosts > 0 and self.launcher is not None - host_mesh_extent = host_mesh.extent if host_mesh else None + is_remote = ( + num_hosts is not None and num_hosts > 0 and self.launcher is not None + ) + host_mesh_extent = host_mesh.extent if host_mesh else None async with self._lock: server_name = None if is_remote: @@ -340,9 +340,11 @@ async def get_proc_mesh( num_hosts=num_hosts, ) - assert host_mesh, f"host_mesh cannot be None at this point." + assert ( + host_mesh + ), f"{mesh_name=}: host_mesh cannot be None at this point." await self.add_host_mesh(host_mesh, mesh_name) - + host_id = host_mesh._host_id gpu_manager = self._host_gpu_map[host_id] else: diff --git a/src/forge/types.py b/src/forge/types.py index 3d91bd502..ebc13cd76 100644 --- a/src/forge/types.py +++ b/src/forge/types.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from enum import Enum from typing import Any, TypedDict, Union + from omegaconf import DictConfig @@ -119,14 +120,18 @@ class LauncherConfig: account: str = "" qos: str = "" ssh_hostfile: str = None - colocate_ref_and_trainer: bool = True + colocate_ref_and_trainer: bool = True def __post_init__(self): if isinstance(self.launcher, str): self.launcher = Launcher(self.launcher) - self.services = {name: ServiceConfig(**value) for name, value in self.services.items()} - self.actors = {name: ProcessConfig(**value) for name, value in self.actors.items()} + self.services = { + name: ServiceConfig(**value) for name, value in self.services.items() + } + self.actors = { + name: ProcessConfig(**value) for name, value in self.actors.items() + } @classmethod def prepare_kwargs(cls, cfg: DictConfig) -> dict: @@ -136,6 +141,7 @@ def prepare_kwargs(cls, cfg: DictConfig) -> dict: kwargs["actors"] = cfg.actors return kwargs + @dataclass class ProvisionerConfig: """A config for the forge provisioner.""" From e98cc4e1d81a2fc3114b2fd4086d54d374ab5f1b Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Tue, 30 Dec 2025 18:06:48 +0000 Subject: [PATCH 03/11] Revert changes Signed-off-by: Olatunji Ruwase --- apps/grpo/qwen3_1_7b.yaml | 43 +++++++------------ .../controller/launchers/ssh_launcher.py | 7 --- 2 files changed, 16 insertions(+), 34 deletions(-) diff --git a/apps/grpo/qwen3_1_7b.yaml b/apps/grpo/qwen3_1_7b.yaml index a9bb5b85e..dc21d2f72 100644 --- a/apps/grpo/qwen3_1_7b.yaml +++ b/apps/grpo/qwen3_1_7b.yaml @@ -10,19 +10,12 @@ model: "Qwen/Qwen3-1.7B" off_by_n: 1 # Off by one by default compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM -# Main loop configuration -rollout_threads: ${services.policy.num_replicas} # Recommended to set equal to policy.num_replicas - -provisioner: - launcher: ssh - ssh_hostfile: "/job/hostfile" - # Observability configuration metric_logging: - # wandb: - # project: grpo-training - # group: grpo_exp_${oc.env:USER} - # logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce + wandb: + project: grpo-training + group: grpo_exp_${oc.env:USER} + logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce console: logging_mode: global_reduce @@ -70,7 +63,7 @@ trainer: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: 4 + data_parallel_shard_degree: -1 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -78,22 +71,21 @@ trainer: disable_loss_parallel: true checkpoint: enable: true - folder: /data-fast/torchforge/${trainer.model.flavor} - initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. - initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo - last_save_in_hf: true - interval: 500 - async_mode: "disabled" + folder: ./checkpoint # The folder to save checkpoints to. + initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. + initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo + last_save_in_hf: true + interval: 500 + async_mode: "disabled" activation_checkpoint: mode: selective selective_ac_option: op - dcp_path: /modeling-checkpoints/torchforge_dcp_tmp/${trainer.model.flavor} # Replay buffer configuration replay_buffer: batch_size: ${local_batch_size} max_policy_age: ${off_by_n} - dp_size: ${trainer.parallelism.data_parallel_shard_degree} # Must equal trainer DP degree + dp_size: 2 # Reference model configuration ref_model: @@ -123,16 +115,14 @@ ref_model: services: policy: procs: ${policy.engine_args.tensor_parallel_size} - num_replicas: 4 - mesh_name: policy + num_replicas: 1 with_gpus: true - hosts: 1 + mesh_name: policy ref_model: procs: 1 - num_replicas: 4 + num_replicas: 1 mesh_name: ref_model with_gpus: true - hosts: 1 reward_actor: procs: 1 num_replicas: 1 @@ -145,10 +135,9 @@ actors: with_gpus: false mesh_name: dataset trainer: - procs: 4 + procs: 2 with_gpus: true mesh_name: trainer - hosts: 1 replay_buffer: procs: 1 with_gpus: false diff --git a/src/forge/controller/launchers/ssh_launcher.py b/src/forge/controller/launchers/ssh_launcher.py index f4ed7c046..4731eeb51 100644 --- a/src/forge/controller/launchers/ssh_launcher.py +++ b/src/forge/controller/launchers/ssh_launcher.py @@ -161,7 +161,6 @@ async def _setup_job(self): return job async def get_host_mesh(self, name: str, num_hosts: int) -> tuple[HostMesh, str]: - print(f"ssh_launcher_get_host_mesh: {name=} {num_hosts=}") for key in self.host_mesh_map.keys(): if key in name: host_mesh = self.host_mesh_map[key].get_host_mesh() @@ -213,10 +212,6 @@ def validate_configuration(cls, cfg: LauncherConfig) -> tuple[bool, str]: return False, error_msg # Check that all services/actors which use gpus also specify host count - print(f"{cfg=}") - print(f"{cfg.services=}") - print(f"{cfg.actors=}") - gpu_host_count = { name: service.hosts for name, service in cfg.services.items() @@ -250,6 +245,4 @@ def validate_configuration(cls, cfg: LauncherConfig) -> tuple[bool, str]: error_msg = f"{error_msg_prefix} {required_host_count=} exceeds {available_host_count=}" return False, error_msg - print(f"{required_host_count=} {available_host_count=} {gpu_host_count=}") - return True, None From 16512ea2810e4989a7eebe7d012c42a518308850 Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Tue, 30 Dec 2025 18:13:39 +0000 Subject: [PATCH 04/11] Revert yaml changes Signed-off-by: Olatunji Ruwase --- apps/grpo/llama3_8b.yaml | 36 ++++++++++++------------------------ apps/grpo/qwen3_1_7b.yaml | 6 +++--- apps/grpo/qwen3_8b.yaml | 39 +++++++++++++-------------------------- 3 files changed, 28 insertions(+), 53 deletions(-) diff --git a/apps/grpo/llama3_8b.yaml b/apps/grpo/llama3_8b.yaml index 32062c8a2..6a887ebc3 100644 --- a/apps/grpo/llama3_8b.yaml +++ b/apps/grpo/llama3_8b.yaml @@ -10,19 +10,12 @@ model: "meta-llama/Meta-Llama-3.1-8B-Instruct" off_by_n: 1 # Off by one by default compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM -# Main loop configuration -rollout_threads: ${services.policy.num_replicas} # Recommended to set equal to policy.num_replicas - -provisioner: - launcher: ssh - ssh_hostfile: "/job/hostfile" - # Observability configuration metric_logging: - # wandb: - # project: grpo-training - # group: grpo_exp_${oc.env:USER} - # logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce + wandb: + project: grpo-training + group: grpo_exp_${oc.env:USER} + logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce console: logging_mode: global_reduce @@ -70,7 +63,7 @@ trainer: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: 4 + data_parallel_shard_degree: -1 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -78,7 +71,7 @@ trainer: disable_loss_parallel: true checkpoint: enable: true - folder: /data-fast/torchforge/${trainer.model.flavor} + folder: ./checkpoint # The folder to save checkpoints to. initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo last_save_in_hf: true @@ -87,7 +80,6 @@ trainer: activation_checkpoint: mode: selective selective_ac_option: op - dcp_path: /modeling-checkpoints/torchforge_dcp_tmp/${trainer.model.flavor} # Replay buffer configuration replay_buffer: @@ -96,7 +88,7 @@ replay_buffer: # This should match the dp_size of TorchTitan # Here it's set explicitly to 2, because we've set # 2 GPUs for the trainer and we're using full FSDP. - dp_size: ${trainer.parallelism.data_parallel_shard_degree} # Must equal trainer DP degree + dp_size: 2 # Reference model configuration ref_model: @@ -112,7 +104,7 @@ ref_model: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: 4 + data_parallel_shard_degree: 1 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -125,17 +117,14 @@ ref_model: services: policy: procs: ${policy.engine_args.tensor_parallel_size} - num_replicas: 2 - mesh_name: policy + num_replicas: 1 with_gpus: true - hosts: 1 - + mesh_name: policy ref_model: - procs: 4 + procs: 1 num_replicas: 1 with_gpus: true mesh_name: ref_model - hosts: 1 reward_actor: procs: 1 num_replicas: 1 @@ -148,10 +137,9 @@ actors: with_gpus: false mesh_name: dataset trainer: - procs: 4 + procs: 2 with_gpus: true mesh_name: trainer - hosts: 1 replay_buffer: procs: 1 with_gpus: false diff --git a/apps/grpo/qwen3_1_7b.yaml b/apps/grpo/qwen3_1_7b.yaml index dc21d2f72..45cdcbf74 100644 --- a/apps/grpo/qwen3_1_7b.yaml +++ b/apps/grpo/qwen3_1_7b.yaml @@ -85,7 +85,7 @@ trainer: replay_buffer: batch_size: ${local_batch_size} max_policy_age: ${off_by_n} - dp_size: 2 + dp_size: ${trainer.parallelism.data_parallel_shard_degree} # Must equal trainer DP degree # Reference model configuration ref_model: @@ -116,8 +116,8 @@ services: policy: procs: ${policy.engine_args.tensor_parallel_size} num_replicas: 1 - with_gpus: true mesh_name: policy + with_gpus: true ref_model: procs: 1 num_replicas: 1 @@ -135,7 +135,7 @@ actors: with_gpus: false mesh_name: dataset trainer: - procs: 2 + procs: 1 with_gpus: true mesh_name: trainer replay_buffer: diff --git a/apps/grpo/qwen3_8b.yaml b/apps/grpo/qwen3_8b.yaml index 50807ae58..9bec6a541 100644 --- a/apps/grpo/qwen3_8b.yaml +++ b/apps/grpo/qwen3_8b.yaml @@ -10,20 +10,12 @@ model: "Qwen/Qwen3-8B" off_by_n: 1 # Off by one by default compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM -# Main loop configuration -rollout_threads: ${services.policy.num_replicas} # Recommended to set equal to policy.num_replicas - -provisioner: - launcher: ssh - ssh_hostfile: "/job/hostfile" - colocate_ref_and_trainer: true - # Observability configuration metric_logging: - # wandb: - # project: grpo-training - # group: grpo_exp_${oc.env:USER} - # logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce + wandb: + project: grpo-training + group: grpo_exp_${oc.env:USER} + logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce console: logging_mode: global_reduce @@ -41,7 +33,7 @@ policy: model: ${model} tensor_parallel_size: 2 pipeline_parallel_size: 1 - enforce_eager: true # ${not:${compile}} + enforce_eager: ${not:${compile}} sampling_params: # https://docs.vllm.ai/en/v0.10.0/api/vllm/sampling_params.html#vllm.sampling_params.SamplingParams n: ${group_size} max_tokens: ${max_res_tokens} @@ -71,7 +63,7 @@ trainer: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: 4 + data_parallel_shard_degree: -1 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -79,7 +71,7 @@ trainer: disable_loss_parallel: true checkpoint: enable: true - folder: /data-fast/torchforge/${trainer.model.flavor} + folder: ./checkpoint # The folder to save checkpoints to. initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo last_save_in_hf: true @@ -88,7 +80,6 @@ trainer: activation_checkpoint: mode: selective selective_ac_option: op - dcp_path: /modeling-checkpoints/torchforge_dcp_tmp/${trainer.model.flavor} # Replay buffer configuration replay_buffer: @@ -97,7 +88,7 @@ replay_buffer: # This should match the dp_size of TorchTitan # Here it's set explicitly to 2, because we've set # 2 GPUs for the trainer and we're using full FSDP. - dp_size: ${trainer.parallelism.data_parallel_shard_degree} # Must equal trainer DP degree + dp_size: 2 # Reference model configuration ref_model: @@ -113,7 +104,7 @@ ref_model: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: 4 + data_parallel_shard_degree: 1 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -126,17 +117,14 @@ ref_model: services: policy: procs: ${policy.engine_args.tensor_parallel_size} - num_replicas: 2 - mesh_name: policy + num_replicas: 1 with_gpus: true - hosts: 1 - + mesh_name: policy ref_model: - procs: 4 + procs: 1 num_replicas: 1 with_gpus: true mesh_name: ref_model - hosts: 1 reward_actor: procs: 1 num_replicas: 1 @@ -149,10 +137,9 @@ actors: with_gpus: false mesh_name: dataset trainer: - procs: 4 + procs: 2 with_gpus: true mesh_name: trainer - hosts: 1 replay_buffer: procs: 1 with_gpus: false From 3c3c8c481baa8a9ebeedcea771e56a80fd76651a Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Tue, 30 Dec 2025 18:15:06 +0000 Subject: [PATCH 05/11] Revert yaml Signed-off-by: Olatunji Ruwase --- apps/grpo/qwen3_1_7b.yaml | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/apps/grpo/qwen3_1_7b.yaml b/apps/grpo/qwen3_1_7b.yaml index 45cdcbf74..b0c08a28c 100644 --- a/apps/grpo/qwen3_1_7b.yaml +++ b/apps/grpo/qwen3_1_7b.yaml @@ -10,6 +10,10 @@ model: "Qwen/Qwen3-1.7B" off_by_n: 1 # Off by one by default compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM +# Main loop configuration +rollout_threads: 1 # Recommended to set equal to policy.num_replicas + + # Observability configuration metric_logging: wandb: @@ -63,7 +67,7 @@ trainer: enable: ${compile} parallelism: data_parallel_replicate_degree: 1 - data_parallel_shard_degree: -1 + data_parallel_shard_degree: 1 tensor_parallel_degree: 1 pipeline_parallel_degree: 1 context_parallel_degree: 1 @@ -71,12 +75,12 @@ trainer: disable_loss_parallel: true checkpoint: enable: true - folder: ./checkpoint # The folder to save checkpoints to. - initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. - initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo - last_save_in_hf: true - interval: 500 - async_mode: "disabled" + folder: ./checkpoint # The folder to save checkpoints to. + initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. + initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo + last_save_in_hf: true + interval: 500 + async_mode: "disabled" activation_checkpoint: mode: selective selective_ac_option: op From 5ca7c628778b75ae2dabc482654d7112f086753c Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Thu, 1 Jan 2026 17:44:58 +0000 Subject: [PATCH 06/11] Cleanup Signed-off-by: Olatunji Ruwase --- src/forge/controller/launchers/launcher_factory.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/forge/controller/launchers/launcher_factory.py b/src/forge/controller/launchers/launcher_factory.py index 47711c875..45cbbd8f0 100644 --- a/src/forge/controller/launchers/launcher_factory.py +++ b/src/forge/controller/launchers/launcher_factory.py @@ -26,10 +26,6 @@ def get_launcher(cfg: LauncherConfig | None = None) -> BaseLauncher | None: success, error_msg = SSHLauncher.validate_configuration(cfg) if not success: raise ValueError(error_msg) - # if cfg.ssh_hostfile is None or not os.path.isfile(cfg.ssh_hostfile): - # raise ValueError( - # f"SSH Launcher cannot work because of invalid hostfile path: {cfg.ssh_hostfile}." - # ) return SSHLauncher(cfg) elif cfg.launcher == Launcher.MAST: try: From 9fe309bdd938b4a2e5f98b6ab6fa1d9e5687be27 Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Fri, 23 Jan 2026 16:35:45 +0000 Subject: [PATCH 07/11] Formatting Signed-off-by: Olatunji Ruwase --- experimental/ssh/qwen3_1_7b.yaml | 164 ++++++++++++++++++ experimental/ssh/qwen3_8b.yaml | 163 +++++++++++++++++ src/forge/controller/launchers/launcher.py | 28 --- .../controller/launchers/slurm_launcher.py | 34 +--- .../controller/launchers/ssh_launcher.py | 152 +++++++--------- src/forge/types.py | 27 ++- 6 files changed, 417 insertions(+), 151 deletions(-) create mode 100644 experimental/ssh/qwen3_1_7b.yaml create mode 100644 experimental/ssh/qwen3_8b.yaml diff --git a/experimental/ssh/qwen3_1_7b.yaml b/experimental/ssh/qwen3_1_7b.yaml new file mode 100644 index 000000000..30a1a1c96 --- /dev/null +++ b/experimental/ssh/qwen3_1_7b.yaml @@ -0,0 +1,164 @@ +# Grouped Relative Policy Optimization (GRPO) +# >>> python -m apps.grpo.main --config experimental/ssh/qwen3_1.7B.yaml + +# Global configuration +group_size: 8 +local_batch_size: 16 # per-device batch size +max_req_tokens: 1024 +max_res_tokens: 2048 +model: "Qwen/Qwen3-1.7B" +off_by_n: 1 # Off by one by default +compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM + +# Main loop configuration +rollout_threads: ${services.generator.num_replicas} # Recommended to set equal to generator.num_replicas + +provisioner: + launcher: ssh + monarch_port: 22222 + ssh_hostfile: /job/hostfile + colocate: [ref_model, trainer] + +# Observability configuration +metric_logging: + wandb: + project: torch_forge_grpo_gsm8k + group: qwen_1.7b + name: ssh_mesh + logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce + console: + logging_mode: global_reduce + +# Dataset configuration +dataset: + path: "openai/gsm8k" + revision: "main" + data_split: "train" + streaming: true + model: ${model} + +# Generator configuration +generator: + engine_args: # https://docs.vllm.ai/en/v0.10.0/api/vllm/engine/arg_utils.html#vllm.engine.arg_utils.EngineArgs + model: ${model} + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + enforce_eager: true # ${not:${compile}} + sampling_params: # https://docs.vllm.ai/en/v0.10.0/api/vllm/sampling_params.html#vllm.sampling_params.SamplingParams + n: ${group_size} + max_tokens: ${max_res_tokens} + temperature: 1.0 + top_p: 1.0 + +# Trainer configuration +trainer: + model: + name: qwen3 + flavor: 1.7B + hf_assets_path: hf://${model} + optimizer: + name: AdamW + lr: 1e-5 + eps: 1e-8 + lr_scheduler: + warmup_steps: 1 + training: + local_batch_size: ${local_batch_size} + seq_len: ${sum:${max_req_tokens},${max_res_tokens}} # seq_len >= max_req_tokens + max_res_tokens + max_norm: 1.0 + steps: 10 # 1000000 + dtype: bfloat16 + gc_freq: 1 + compile: + enable: ${compile} + parallelism: + data_parallel_replicate_degree: 2 + data_parallel_shard_degree: 1 + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + context_parallel_degree: 1 + expert_parallel_degree: 1 + disable_loss_parallel: true + checkpoint: + enable: true + folder: /data-fast/torchforge/${trainer.model.flavor} + initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. + initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo + last_save_in_hf: true + interval: 500 + async_mode: "disabled" + activation_checkpoint: + mode: selective + selective_ac_option: op + +# Replay buffer configuration +replay_buffer: + batch_size: ${local_batch_size} + max_policy_age: ${off_by_n} + # This should match the dp_size of TorchTitan + # Here it's set explicitly to 2, because we've set + # 2 GPUs for the trainer and we're using full FSDP. + dp_size: ${trainer.parallelism.data_parallel_replicate_degree} # Must equal trainer DP degree + +# Reference model configuration +ref_model: + model: + name: qwen3 + flavor: 1.7B + hf_assets_path: hf://${model} + training: + seq_len: ${trainer.training.seq_len} + dtype: bfloat16 + gc_freq: 1 + compile: + enable: ${compile} + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: 1 + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + context_parallel_degree: 1 + expert_parallel_degree: 1 + checkpoint: + initial_load_path: hf://${model} + initial_load_in_hf: true + +# All resource allocations +services: + generator: + procs: ${generator.engine_args.tensor_parallel_size} + num_replicas: 2 + mesh_name: generator + with_gpus: true + hosts: 1 + + ref_model: + procs: 1 + num_replicas: 2 + with_gpus: true + mesh_name: ref_model + hosts: 1 + reward_actor: + procs: 1 + num_replicas: 1 + with_gpus: false + mesh_name: reward_actor + +actors: + dataset: + procs: 1 + with_gpus: false + mesh_name: dataset + trainer: + procs: 2 + with_gpus: true + mesh_name: trainer + hosts: 1 + replay_buffer: + procs: 1 + with_gpus: false + mesh_name: replay_buffer + compute_advantages: + procs: 1 + with_gpus: false + mesh_name: compute_advantages diff --git a/experimental/ssh/qwen3_8b.yaml b/experimental/ssh/qwen3_8b.yaml new file mode 100644 index 000000000..616593eeb --- /dev/null +++ b/experimental/ssh/qwen3_8b.yaml @@ -0,0 +1,163 @@ +# Grouped Relative Policy Optimization (GRPO) +# >>> python -m apps.grpo.main --config experimental/ssh/qwen3_8b.yaml + +# Global configuration +group_size: 16 +local_batch_size: 4 # per-device batch size +max_req_tokens: 1024 +max_res_tokens: 2048 +model: "Qwen/Qwen3-8B" +off_by_n: 1 # Off by one by default +compile: true # Enable torch.compile for trainer/ref_model, and CUDA graphs for vLLM + +# Main loop configuration +rollout_threads: ${services.generator.num_replicas} # Recommended to set equal to generator.num_replicas + +provisioner: + launcher: ssh + monarch_port: 22222 + ssh_hostfile: /job/hostfile + colocate: [ref_model, trainer] + +# Observability configuration +metric_logging: + wandb: + project: torch_forge_grpo_gsm8k + group: qwen3_8b + name: ssh_mesh + logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce + console: + logging_mode: global_reduce + +# Dataset configuration +dataset: + path: "openai/gsm8k" + revision: "main" + data_split: "train" + streaming: true + model: ${model} + +# Generator configuration +generator: + engine_args: # https://docs.vllm.ai/en/v0.10.0/api/vllm/engine/arg_utils.html#vllm.engine.arg_utils.EngineArgs + model: ${model} + tensor_parallel_size: 2 + pipeline_parallel_size: 1 + enforce_eager: true # ${not:${compile}} + sampling_params: # https://docs.vllm.ai/en/v0.10.0/api/vllm/sampling_params.html#vllm.sampling_params.SamplingParams + n: ${group_size} + max_tokens: ${max_res_tokens} + temperature: 1.0 + top_p: 1.0 + +# Trainer configuration +trainer: + model: + name: qwen3 + flavor: 8B + hf_assets_path: hf://${model} + optimizer: + name: AdamW + lr: 1e-5 + eps: 1e-8 + lr_scheduler: + warmup_steps: 1 + training: + local_batch_size: ${local_batch_size} + seq_len: ${sum:${max_req_tokens},${max_res_tokens}} # seq_len >= max_req_tokens + max_res_tokens + max_norm: 1.0 + steps: 10 + dtype: bfloat16 + gc_freq: 1 + compile: + enable: ${compile} + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: 4 + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + context_parallel_degree: 1 + expert_parallel_degree: 1 + disable_loss_parallel: true + checkpoint: + enable: true + folder: /data-fast/torchforge/${trainer.model.flavor} + initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. + initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo + last_save_in_hf: true + interval: 500 + async_mode: "disabled" + activation_checkpoint: + mode: selective + selective_ac_option: op + +# Replay buffer configuration +replay_buffer: + batch_size: ${local_batch_size} + max_policy_age: ${off_by_n} + # This should match the dp_size of TorchTitan + # Here it's set explicitly to 2, because we've set + # 2 GPUs for the trainer and we're using full FSDP. + dp_size: ${trainer.parallelism.data_parallel_shard_degree} # Must equal trainer DP degree + +# Reference model configuration +ref_model: + model: + name: qwen3 + flavor: 8B + hf_assets_path: hf://${model} + training: + seq_len: ${trainer.training.seq_len} + dtype: bfloat16 + gc_freq: 1 + compile: + enable: ${compile} + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: 2 + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + context_parallel_degree: 1 + expert_parallel_degree: 1 + checkpoint: + initial_load_path: hf://${model} + initial_load_in_hf: true + +# All resource allocations +services: + generator: + procs: ${generator.engine_args.tensor_parallel_size} + num_replicas: 2 + hosts: 1 + with_gpus: true + mesh_name: generator + ref_model: + procs: 2 + num_replicas: 2 + hosts: 1 + with_gpus: true + mesh_name: ref_model + reward_actor: + procs: 1 + num_replicas: 1 + with_gpus: false + mesh_name: reward_actor + +actors: + dataset: + procs: 1 + with_gpus: false + mesh_name: dataset + trainer: + procs: 4 + hosts: 1 + with_gpus: true + mesh_name: trainer + replay_buffer: + procs: 1 + with_gpus: false + mesh_name: replay_buffer + compute_advantages: + procs: 1 + with_gpus: false + mesh_name: compute_advantages diff --git a/src/forge/controller/launchers/launcher.py b/src/forge/controller/launchers/launcher.py index 8905a1712..508e39ffd 100644 --- a/src/forge/controller/launchers/launcher.py +++ b/src/forge/controller/launchers/launcher.py @@ -17,34 +17,6 @@ LAUNCHER_KEY = "launcher" -def get_meshes_from_config(cfg: LauncherConfig) -> dict[str, int]: - """Extract mesh requirements from launcher config. - - Args: - cfg: The launcher configuration - - Returns: - Dictionary mapping mesh names to number of hosts required - """ - meshes: dict[str, int] = {} - - # Add services that need remote hosts - for service_name, service_cfg in cfg.services.items(): - hosts = getattr(service_cfg, "hosts", None) - if hosts and hosts > 0: - mesh_name = service_cfg.mesh_name or service_name - meshes[mesh_name] = hosts - - # Add actors that need remote hosts - for actor_name, actor_cfg in cfg.actors.items(): - hosts = getattr(actor_cfg, "hosts", None) - if hosts and hosts > 0: - mesh_name = actor_cfg.mesh_name or actor_name - meshes[mesh_name] = hosts - - return meshes - - def get_launcher(cfg: LauncherConfig | None = None) -> BaseLauncher | None: if not cfg: return None diff --git a/src/forge/controller/launchers/slurm_launcher.py b/src/forge/controller/launchers/slurm_launcher.py index b0ec7a10d..1bfa2203e 100644 --- a/src/forge/controller/launchers/slurm_launcher.py +++ b/src/forge/controller/launchers/slurm_launcher.py @@ -20,38 +20,6 @@ logger.setLevel(logging.DEBUG) -JOB_NAME_KEY = "job_name" -LAUNCHER_KEY = "launcher" - - -def get_meshes_from_config(cfg: LauncherConfig) -> dict[str, int]: - """Extract mesh requirements from launcher config. - - Args: - cfg: The launcher configuration - - Returns: - Dictionary mapping mesh names to number of hosts required - """ - meshes: dict[str, int] = {} - - # Add services that need remote hosts - for service_name, service_cfg in cfg.services.items(): - hosts = getattr(service_cfg, "hosts", None) - if hosts and hosts > 0: - mesh_name = service_cfg.mesh_name or service_name - meshes[mesh_name] = hosts - - # Add actors that need remote hosts - for actor_name, actor_cfg in cfg.actors.items(): - hosts = getattr(actor_cfg, "hosts", None) - if hosts and hosts > 0: - mesh_name = actor_cfg.mesh_name or actor_name - meshes[mesh_name] = hosts - - return meshes - - class Slurmlauncher(BaseLauncher): def __init__( self, @@ -68,7 +36,7 @@ async def initialize(self) -> tuple[JobTrait, JobState]: A tuple of (job, job_state) containing the SlurmJob handle and its state. """ # Collect all mesh requirements from config - meshes = get_meshes_from_config(self.cfg) + meshes = self.cfg.get_meshes() # If no remote resources needed, skip job creation if not meshes: diff --git a/src/forge/controller/launchers/ssh_launcher.py b/src/forge/controller/launchers/ssh_launcher.py index 520b0b1f9..9ad6955cb 100644 --- a/src/forge/controller/launchers/ssh_launcher.py +++ b/src/forge/controller/launchers/ssh_launcher.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""SSHLauncher """ +"""SSHLauncher""" # from .launcher import get_meshes_from_config import atexit @@ -12,7 +12,6 @@ import os import re import sys - from dataclasses import dataclass, field from typing import List @@ -68,15 +67,10 @@ def parse_hostfile(hostfile_lines): return resource_pool -GENERATOR_KEY = "generator" -REFERENCE_KEY = "ref_model" -TRAINER_KEY = "trainer" -ProcMesh_KEYS = [GENERATOR_KEY, REFERENCE_KEY, TRAINER_KEY] - - @dataclass class SSHForgeProc: is_actor: bool = True + is_colocate: bool = False num_hosts: int = 0 @@ -101,34 +95,6 @@ def get_host_mesh(self): return host_mesh -def get_meshes_from_config(cfg: LauncherConfig) -> dict[str, int]: - """Extract mesh requirements from launcher config. - - Args: - cfg: The launcher configuration - - Returns: - Dictionary mapping mesh names to number of hosts required - """ - meshes: dict[str, int] = {} - - # Add services that need remote hosts - for service_name, service_cfg in cfg.services.items(): - hosts = getattr(service_cfg, "hosts", None) - if hosts and hosts > 0: - mesh_name = service_cfg.mesh_name or service_name - meshes[mesh_name] = hosts - - # Add actors that need remote hosts - for actor_name, actor_cfg in cfg.actors.items(): - hosts = getattr(actor_cfg, "hosts", None) - if hosts and hosts > 0: - mesh_name = actor_cfg.mesh_name or actor_name - meshes[mesh_name] = hosts - - return meshes - - class SSHLauncher(BaseLauncher): def __init__(self, cfg: LauncherConfig): self.cfg = cfg @@ -137,42 +103,57 @@ def __init__(self, cfg: LauncherConfig): self.ssh_hostfile = cfg.ssh_hostfile self.host_pool = fetch_hostfile(self.ssh_hostfile) self.host_ips = list(self.host_pool.keys()) + self.monarch_port = cfg.monarch_port - self.meshes = get_meshes_from_config(self.cfg) - print(f"sshlauncher: {self.meshes=}") - + self.proc_meshes = self.cfg.get_meshes() + self.colocate_procs = [p for p in self.cfg.colocate if p in self.proc_meshes] self.forge_proc_map = self._create_forge_proc_map(cfg) - n_generator_slice_end = self.forge_proc_map[GENERATOR_KEY].num_hosts - n_trainer_slice_end = ( - n_generator_slice_end + self.forge_proc_map[TRAINER_KEY].num_hosts - ) - n_ref_model_slice_start = ( - n_generator_slice_end - if cfg.colocate_ref_and_trainer - else n_trainer_slice_end - ) - n_ref_model_slice_end = ( - n_ref_model_slice_start + self.forge_proc_map[REFERENCE_KEY].num_hosts - ) - self.host_ip_map = { - GENERATOR_KEY: self.host_ips[:n_generator_slice_end], - TRAINER_KEY: self.host_ips[n_generator_slice_end:n_trainer_slice_end], - REFERENCE_KEY: self.host_ips[n_ref_model_slice_start:n_ref_model_slice_end], - } + required_host_count = sum( + [ + proc.num_hosts + for proc in self.forge_proc_map.values() + if not proc.is_colocate + ] + ) + colocated_host_count = 0 + if len(self.colocate_procs) > 0: + colocated_host_count = max( + [ + proc.num_hosts + for proc in self.forge_proc_map.values() + if proc.is_colocate + ] + ) + required_host_count += colocated_host_count + + self.host_ip_map = {} + slice_start = 0 + for key, proc in self.forge_proc_map.items(): + if proc.is_colocate: + # colocated procs share the last block of hosts + self.host_ip_map[key] = self.host_ips[-colocated_host_count:] + else: + self.host_ip_map[key] = self.host_ips[slice_start : proc.num_hosts] + slice_start += proc.num_hosts self.host_mesh_map = dict() def _create_forge_proc_map(self, cfg: LauncherConfig) -> dict: proc_map = {} - for key in ProcMesh_KEYS: # self.meshes: + for key in self.proc_meshes.keys(): + is_colocate = key in self.colocate_procs if key in cfg.services: proc_map[key] = SSHForgeProc( - is_actor=False, num_hosts=cfg.services[key].hosts + is_actor=False, + is_colocate=is_colocate, + num_hosts=cfg.services[key].hosts, ) else: proc_map[key] = SSHForgeProc( - is_actor=True, num_hosts=cfg.actors[key].hosts + is_actor=True, + is_colocate=is_colocate, + num_hosts=cfg.actors[key].hosts, ) return proc_map @@ -186,14 +167,10 @@ async def _setup_job(self): "-o", "BatchMode=yes", ], - monarch_port=int( - os.getenv("MONARCH_PORT", "22222") - ), # make sure this is open + monarch_port=self.monarch_port, # make sure this is open ) - # for proc_key in self.meshes: - for proc_key in ProcMesh_KEYS: - # job.add_mesh(f"{proc_key}_mesh", self.host_ip_map[proc_key]) + for proc_key in self.proc_meshes.keys(): job.add_mesh(f"{proc_key}", self.host_ip_map[proc_key]) return job @@ -227,21 +204,10 @@ async def initialize(self) -> None: self.job = await self._setup_job() self.state = self.job.state() - self.host_mesh_map[TRAINER_KEY] = self._create_actor_or_service( - mesh_name=TRAINER_KEY, - job_host_mesh=self.state.trainer, - # job_host_mesh=self.state.trainer_mesh, - ) - self.host_mesh_map[GENERATOR_KEY] = self._create_actor_or_service( - mesh_name=GENERATOR_KEY, - job_host_mesh=self.state.generator, - # job_host_mesh=self.state.generator_mesh, - ) - self.host_mesh_map[REFERENCE_KEY] = self._create_actor_or_service( - mesh_name=REFERENCE_KEY, - job_host_mesh=self.state.ref_model, - # job_host_mesh=self.state.ref_model_mesh, - ) + for proc_key in self.forge_proc_map.keys(): + self.host_mesh_map[proc_key] = self._create_actor_or_service( + mesh_name=proc_key, job_host_mesh=getattr(self.state, proc_key) + ) # Register cleanup handler atexit.register(self.job.kill) @@ -280,17 +246,25 @@ def validate_configuration(cls, cfg: LauncherConfig) -> tuple[bool, str]: # Check that sufficient hosts available available_host_count = len(fetch_hostfile(cfg.ssh_hostfile)) + colocate_hosts = [ + host for host in gpu_host_count.keys() if host in cfg.colocate + ] required_host_count = sum(list(gpu_host_count.values())) - - if ( - cfg.colocate_ref_and_trainer - and TRAINER_KEY in gpu_host_count - and REFERENCE_KEY in gpu_host_count - ): - excess_gpu_count = min( - gpu_host_count[TRAINER_KEY], gpu_host_count[REFERENCE_KEY] + required_host_count = sum( + [ + count + for name, count in gpu_host_count.items() + if name not in cfg.colocate + ] + ) + if len(colocate_hosts) > 0: + required_host_count += max( + [ + count + for name, count in gpu_host_count.items() + if name in cfg.colocate + ] ) - required_host_count -= excess_gpu_count if required_host_count > available_host_count: error_msg = f"{error_msg_prefix} {required_host_count=} exceeds {available_host_count=}" diff --git a/src/forge/types.py b/src/forge/types.py index 0d656d2c4..d54a7f772 100644 --- a/src/forge/types.py +++ b/src/forge/types.py @@ -116,13 +116,38 @@ class LauncherConfig: None # Memory per node (SLURM param, can get with sinfo) ) ssh_hostfile: str = None - colocate_ref_and_trainer: bool = True + monarch_port: int = 22222 + colocate: list[str] = field(default_factory=list) # services and actors to colocate gpus_per_node: int = 8 # GPUs per node (SLURM param, can get with sinfo) def __post_init__(self): if isinstance(self.launcher, str): self.launcher = Launcher(self.launcher) + def get_meshes(self) -> dict[str, int]: + """Extract mesh requirements from launcher config. + + Returns: + Dictionary mapping mesh names to number of hosts required + """ + meshes: dict[str, int] = {} + + # Add services that need remote hosts + for service_name, service_cfg in self.services.items(): + hosts = getattr(service_cfg, "hosts", None) + if hosts and hosts > 0: + mesh_name = service_cfg.mesh_name or service_name + meshes[mesh_name] = hosts + + # Add actors that need remote hosts + for actor_name, actor_cfg in self.actors.items(): + hosts = getattr(actor_cfg, "hosts", None) + if hosts and hosts > 0: + mesh_name = actor_cfg.mesh_name or actor_name + meshes[mesh_name] = hosts + + return meshes + @dataclass class ProvisionerConfig: From d68b246224c8162b8a50902b5f3ee635aa6e7678 Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Fri, 23 Jan 2026 16:43:44 +0000 Subject: [PATCH 08/11] Revert unintended changes Signed-off-by: Olatunji Ruwase --- apps/grpo/qwen3_1_7b.yaml | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/apps/grpo/qwen3_1_7b.yaml b/apps/grpo/qwen3_1_7b.yaml index 8fa6bdf2b..e452bc0b0 100644 --- a/apps/grpo/qwen3_1_7b.yaml +++ b/apps/grpo/qwen3_1_7b.yaml @@ -16,10 +16,10 @@ rollout_threads: 1 # Recommended to set equal to generator.num_replicas # Observability configuration metric_logging: - # wandb: - # project: grpo-training - # group: grpo_exp_${oc.env:USER} - # logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce + wandb: + project: grpo-training + group: grpo_exp_${oc.env:USER} + logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce console: logging_mode: global_reduce @@ -37,7 +37,7 @@ generator: model: ${model} tensor_parallel_size: 1 pipeline_parallel_size: 1 - enforce_eager: true # ${not:${compile}} + enforce_eager: ${not:${compile}} sampling_params: # https://docs.vllm.ai/en/v0.10.0/api/vllm/sampling_params.html#vllm.sampling_params.SamplingParams n: ${group_size} max_tokens: ${max_res_tokens} @@ -60,7 +60,7 @@ trainer: local_batch_size: ${local_batch_size} seq_len: ${sum:${max_req_tokens},${max_res_tokens}} # seq_len >= max_req_tokens + max_res_tokens max_norm: 1.0 - steps: 10 + steps: 1000000 dtype: bfloat16 gc_freq: 1 compile: @@ -75,7 +75,7 @@ trainer: disable_loss_parallel: true checkpoint: enable: true - folder: /data-fast/torchforge/${trainer.model.flavor} + folder: ./checkpoint # The folder to save checkpoints to. initial_load_path: hf://${model} # The path to load the initial checkpoint from. Ignored if `folder` exists. initial_load_in_hf: true # If true, interpret initial_load_path as a HuggingFace model repo last_save_in_hf: true @@ -111,7 +111,7 @@ ref_model: context_parallel_degree: 1 expert_parallel_degree: 1 checkpoint: - # enable: true + enable: true initial_load_path: hf://${model} initial_load_in_hf: true From 564f945a8099064e5921ea1e61eb1c2b11cfdf3e Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Fri, 23 Jan 2026 21:07:29 +0000 Subject: [PATCH 09/11] Update API Signed-off-by: Olatunji Ruwase --- src/forge/controller/launchers/ssh_launcher.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/forge/controller/launchers/ssh_launcher.py b/src/forge/controller/launchers/ssh_launcher.py index 9ad6955cb..4e7f9b702 100644 --- a/src/forge/controller/launchers/ssh_launcher.py +++ b/src/forge/controller/launchers/ssh_launcher.py @@ -16,8 +16,6 @@ from typing import List from forge.types import LauncherConfig -from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport -from monarch._rust_bindings.monarch_hyperactor.config import configure from monarch._src.job.job import SSHJob from monarch.actor import HostMesh, ProcMesh @@ -196,11 +194,6 @@ def _create_actor_or_service(self, mesh_name, job_host_mesh): ) async def initialize(self) -> None: - # HostMesh currently requires explicit configuration - # of the underlying transport from client to mesh. - # This can be removed in the future once this has been removed. - configure(default_transport=ChannelTransport.TcpWithHostname) - self.job = await self._setup_job() self.state = self.job.state() From e5600fa8636289ff5d8880af5737a009c04a545c Mon Sep 17 00:00:00 2001 From: Tunji Ruwase Date: Mon, 2 Mar 2026 15:44:27 +0000 Subject: [PATCH 10/11] Fix indent Signed-off-by: Tunji Ruwase --- src/forge/types.py | 53 +++++++++++++++++++++++----------------------- 1 file changed, 26 insertions(+), 27 deletions(-) diff --git a/src/forge/types.py b/src/forge/types.py index a6b39d707..5acd1b6cc 100644 --- a/src/forge/types.py +++ b/src/forge/types.py @@ -124,37 +124,36 @@ def __post_init__(self): if isinstance(self.launcher, str): self.launcher = Launcher(self.launcher) + def get_meshes(self) -> dict[str, int]: + """Extract mesh requirements from launcher config. -def get_meshes(self) -> dict[str, int]: - """Extract mesh requirements from launcher config. + Args: + cfg: The launcher configuration - Args: - cfg: The launcher configuration - - Returns: - Dictionary mapping mesh names to number of hosts required - """ - meshes: dict[str, int] = {} - - # Add services that need remote hosts - # Expand services with multiple replicas into per-replica meshes - for service_name, service_cfg in cfg.services.items(): - hosts = getattr(service_cfg, "hosts", None) - if hosts and hosts > 0: - base_mesh_name = service_cfg.mesh_name or service_name - num_replicas = service_cfg.num_replicas - for replica_idx in range(num_replicas): - mesh_name = f"{base_mesh_name}_{replica_idx}" + Returns: + Dictionary mapping mesh names to number of hosts required + """ + meshes: dict[str, int] = {} + + # Add services that need remote hosts + # Expand services with multiple replicas into per-replica meshes + for service_name, service_cfg in cfg.services.items(): + hosts = getattr(service_cfg, "hosts", None) + if hosts and hosts > 0: + base_mesh_name = service_cfg.mesh_name or service_name + num_replicas = service_cfg.num_replicas + for replica_idx in range(num_replicas): + mesh_name = f"{base_mesh_name}_{replica_idx}" + meshes[mesh_name] = hosts + + # Add actors that need remote hosts + for actor_name, actor_cfg in cfg.actors.items(): + hosts = getattr(actor_cfg, "hosts", None) + if hosts and hosts > 0: + mesh_name = actor_cfg.mesh_name or actor_name meshes[mesh_name] = hosts - # Add actors that need remote hosts - for actor_name, actor_cfg in cfg.actors.items(): - hosts = getattr(actor_cfg, "hosts", None) - if hosts and hosts > 0: - mesh_name = actor_cfg.mesh_name or actor_name - meshes[mesh_name] = hosts - - return meshes + return meshes @dataclass From 55c3b5aa9510499102917d15d4e59b3913b3cf2e Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Tue, 10 Mar 2026 22:39:08 +0000 Subject: [PATCH 11/11] Optional name Signed-off-by: Olatunji Ruwase --- src/forge/controller/provisioner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/forge/controller/provisioner.py b/src/forge/controller/provisioner.py index 5eac5fc73..e0f9d60fd 100644 --- a/src/forge/controller/provisioner.py +++ b/src/forge/controller/provisioner.py @@ -246,7 +246,7 @@ def get_available_gpus(self) -> list[str]: """Returns a list of available GPU devices.""" return [str(gpu) for gpu in self.available_gpus] - def get_gpus(self, num_gpus: int, mesh_name: str) -> list[str]: + def get_gpus(self, num_gpus: int, mesh_name: str = "anonymous") -> list[str]: """Assigns GPU devices.""" if num_gpus > len(self.available_gpus): raise RuntimeError(