diff --git a/experimental/ssh/qwen3_1_7b.yaml b/experimental/ssh/qwen3_1_7b.yaml new file mode 100644 index 000000000..604e3043a --- /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..b1e41200c --- /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/__init__.py b/src/forge/controller/launchers/__init__.py new file mode 100644 index 000000000..b86dc5c14 --- /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 import BaseLauncher +from .launcher import get_launcher + +__all__ = [ + "BaseLauncher", + "get_launcher", +] diff --git a/src/forge/controller/base.py b/src/forge/controller/launchers/base.py similarity index 100% rename from src/forge/controller/base.py rename to src/forge/controller/launchers/base.py diff --git a/src/forge/controller/launchers/launcher.py b/src/forge/controller/launchers/launcher.py new file mode 100644 index 000000000..508e39ffd --- /dev/null +++ b/src/forge/controller/launchers/launcher.py @@ -0,0 +1,39 @@ +# 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""" + + +from forge.types import Launcher, LauncherConfig + +from .base 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) + 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}") diff --git a/src/forge/controller/launcher.py b/src/forge/controller/launchers/slurm_launcher.py similarity index 54% rename from src/forge/controller/launcher.py rename to src/forge/controller/launchers/slurm_launcher.py index 9b90c42ac..0d4e7a057 100644 --- a/src/forge/controller/launcher.py +++ b/src/forge/controller/launchers/slurm_launcher.py @@ -4,56 +4,22 @@ # 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 atexit import logging -from forge.controller.base import BaseLauncher -from forge.types import Launcher, LauncherConfig +from .base import BaseLauncher +from forge.types import LauncherConfig from monarch.actor import ProcMesh + from monarch.job import JobState, JobTrait, SlurmJob logger = logging.getLogger(__name__) 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 - # 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 - - return meshes - - class Slurmlauncher(BaseLauncher): def __init__( self, @@ -70,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: @@ -108,20 +74,3 @@ async def initialize(self) -> tuple[JobTrait, JobState]: 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) - 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..4e7f9b702 --- /dev/null +++ b/src/forge/controller/launchers/ssh_launcher.py @@ -0,0 +1,266 @@ +# 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""" + +# from .launcher import get_meshes_from_config +import atexit +import collections +import os +import re +import sys +from dataclasses import dataclass, field +from typing import List + +from forge.types import LauncherConfig +from monarch._src.job.job import SSHJob +from monarch.actor import HostMesh, ProcMesh + +from .base import BaseLauncher + + +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 + + +@dataclass +class SSHForgeProc: + is_actor: bool = True + is_colocate: bool = False + 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.cfg = cfg + 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.monarch_port = cfg.monarch_port + + 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) + + 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 self.proc_meshes.keys(): + is_colocate = key in self.colocate_procs + if key in cfg.services: + proc_map[key] = SSHForgeProc( + is_actor=False, + is_colocate=is_colocate, + num_hosts=cfg.services[key].hosts, + ) + else: + proc_map[key] = SSHForgeProc( + is_actor=True, + is_colocate=is_colocate, + 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=self.monarch_port, # make sure this is open + ) + + for proc_key in self.proc_meshes.keys(): + job.add_mesh(f"{proc_key}", self.host_ip_map[proc_key]) + + return job + + async def get_host_mesh(self, name: str, num_hosts: int) -> tuple[HostMesh, str]: + 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: + self.job = await self._setup_job() + self.state = self.job.state() + + 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) + + print(f"ssh_launcher[init]: {self.ssh_hostfile=} {self.host_ips=}") + print(f"ssh_launcher[init]: {self.host_ip_map=}") + return self.job, self.state + + async def remote_setup(self, procs: ProcMesh) -> None: + return + + @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 + 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)) + colocate_hosts = [ + host for host in gpu_host_count.keys() if host in cfg.colocate + ] + required_host_count = sum(list(gpu_host_count.values())) + 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 + ] + ) + + if required_host_count > available_host_count: + error_msg = f"{error_msg_prefix} {required_host_count=} exceeds {available_host_count=}" + return False, error_msg + + return True, None diff --git a/src/forge/controller/provisioner.py b/src/forge/controller/provisioner.py index 235835a15..e0f9d60fd 100644 --- a/src/forge/controller/provisioner.py +++ b/src/forge/controller/provisioner.py @@ -13,7 +13,7 @@ import uuid 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 from monarch._src.actor.actor_mesh import ActorMesh @@ -246,10 +246,13 @@ 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 = "anonymous") -> 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}, " + 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] @@ -330,6 +333,21 @@ async def get_host_mesh(self, name: str) -> HostMesh: return host_mesh + 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, @@ -372,33 +390,35 @@ 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: if is_remote: if host_mesh is None: host_mesh = await self.get_host_mesh( name=mesh_name, ) - 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"{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: # 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 @@ -414,9 +434,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 5c8059d86..5acd1b6cc 100644 --- a/src/forge/types.py +++ b/src/forge/types.py @@ -38,6 +38,7 @@ class Observation: class Launcher(Enum): MAST = "mast" SLURM = "slurm" + SSH = "ssh" @dataclass @@ -114,12 +115,46 @@ class LauncherConfig: mem: int | None = ( # noqa: N815 None # Memory per node (SLURM param, can get with sinfo) ) + ssh_hostfile: str = None + 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. + + 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}" + 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 + @dataclass class ProvisionerConfig: