From fbc299ca6245f8b87fde19451167ba8cd525f075 Mon Sep 17 00:00:00 2001
From: Bihan  Rana <bihan@Bihans-MacBook-Pro.local>
Date: Tue, 4 Aug 2026 22:01:12 +0545
Subject: [PATCH] Support Heterogenous Node Groups

---
 runner/internal/runner/executor/executor.go   |  41 ++++--
 runner/internal/runner/schemas/schemas.go     |   1 +
 .../cli/services/configurators/run.py         |  22 ++--
 src/dstack/_internal/cli/utils/run.py         |   2 +-
 .../_internal/core/backends/slurm/compute.py  |   8 +-
 .../_internal/core/compatibility/runs.py      |  11 ++
 .../_internal/core/models/configurations.py   | 100 ++++++++++++++-
 src/dstack/_internal/core/models/runs.py      |  18 +++
 .../background/pipeline_tasks/jobs_running.py |  84 +++++++++++--
 .../pipeline_tasks/jobs_submitted.py          | 119 +++++++++++++++---
 .../services/jobs/configurators/base.py       |  48 +++++--
 .../server/services/jobs/configurators/dev.py |  10 +-
 .../services/jobs/configurators/service.py    |   5 +-
 .../services/jobs/configurators/task.py       |  44 +++++--
 .../server/services/runs/__init__.py          |   2 +-
 .../_internal/server/services/runs/spec.py    |   2 +-
 src/dstack/_internal/server/testing/common.py |  14 ++-
 src/dstack/_internal/utils/interpolator.py    |  11 +-
 .../_internal/utils/nodes_interpolator.py     |  24 ++++
 .../core/models/test_configurations.py        | 117 +++++++++++++++++
 .../pipeline_tasks/test_node_groups.py        | 115 +++++++++++++++++
 .../pipeline_tasks/test_submitted_jobs.py     | 109 ++++++++++++++--
 .../_internal/server/routers/test_runs.py     |   6 +
 .../services/jobs/configurators/test_task.py  |  62 ++++++++-
 .../_internal/utils/test_interpolator.py      |   5 +
 .../utils/test_nodes_interpolator.py          |  40 ++++++
 26 files changed, 929 insertions(+), 91 deletions(-)
 create mode 100644 src/dstack/_internal/utils/nodes_interpolator.py
 create mode 100644 src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py
 create mode 100644 src/tests/_internal/utils/test_nodes_interpolator.py

diff --git a/runner/internal/runner/executor/executor.go b/runner/internal/runner/executor/executor.go
index 31f3d7fe92..bb86c991bd 100644
--- a/runner/internal/runner/executor/executor.go
+++ b/runner/internal/runner/executor/executor.go
@@ -473,7 +473,15 @@ func (ex *RunExecutor) execJob(ctx context.Context, jobLogFile io.Writer) error
 	nodeRank := ex.jobSpec.JobNum
 	nodesNum := ex.jobSpec.JobsPerReplica
 	gpusPerNodeNum := ex.clusterInfo.GPUSPerJob
-	gpusNum := nodesNum * gpusPerNodeNum
+	gpusNum := 0
+	if len(ex.clusterInfo.GPUSPerNode) > 0 {
+		for _, n := range ex.clusterInfo.GPUSPerNode {
+			gpusNum += n
+		}
+	} else {
+		// Old servers omit gpus_per_node; fall back to homogeneous math.
+		gpusNum = nodesNum * gpusPerNodeNum
+	}
 
 	mpiHostfilePath := filepath.Join(ex.dstackDir, "mpi/hostfile")
 
@@ -544,7 +552,7 @@ func (ex *RunExecutor) execJob(ctx context.Context, jobLogFile io.Writer) error
 		log.Warning(ctx, "failed to include dstack_profile", "path", profilePath, "err", err)
 	}
 
-	if err := writeMpiHostfile(ctx, ex.clusterInfo.JobIPs, gpusPerNodeNum, mpiHostfilePath); err != nil {
+	if err := writeMpiHostfile(ctx, ex.clusterInfo.JobIPs, ex.clusterInfo.GPUSPerNode, gpusPerNodeNum, mpiHostfilePath); err != nil {
 		return fmt.Errorf("write MPI hostfile: %w", err)
 	}
 
@@ -759,7 +767,7 @@ func prepareUserSshDir(user *linuxuser.User) (string, error) {
 	return sshDir, nil
 }
 
-func writeMpiHostfile(ctx context.Context, ips []string, gpusPerNode int, path string) error {
+func writeMpiHostfile(ctx context.Context, ips []string, gpusPerNode []int, fallbackGpusPerJob int, path string) error {
 	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
 		return fmt.Errorf("create MPI hostfile directory: %w", err)
 	}
@@ -775,16 +783,25 @@ func writeMpiHostfile(ctx context.Context, ips []string, gpusPerNode int, path s
 		}
 	}
 	if len(nonEmptyIps) == len(ips) {
-		var template string
-		if gpusPerNode == 0 {
-			// CPU node: the number of slots defaults to the number of processor cores on that host
-			// See: https://docs.open-mpi.org/en/main/launching-apps/scheduling.html#calculating-the-number-of-slots
-			template = "%s\n"
-		} else {
-			template = fmt.Sprintf("%%s slots=%d\n", gpusPerNode)
+		if len(gpusPerNode) > 0 && len(gpusPerNode) != len(ips) {
+			return fmt.Errorf(
+				"gpus_per_node length %d != job_ips length %d",
+				len(gpusPerNode), len(ips),
+			)
 		}
-		for _, ip := range nonEmptyIps {
-			if _, err = fmt.Fprintf(file, template, ip); err != nil {
+		for i, ip := range nonEmptyIps {
+			n := fallbackGpusPerJob
+			if len(gpusPerNode) > 0 {
+				n = gpusPerNode[i]
+			}
+			if n == 0 {
+				// CPU node: the number of slots defaults to the number of processor cores on that host
+				// See: https://docs.open-mpi.org/en/main/launching-apps/scheduling.html#calculating-the-number-of-slots
+				_, err = fmt.Fprintf(file, "%s\n", ip)
+			} else {
+				_, err = fmt.Fprintf(file, "%s slots=%d\n", ip, n)
+			}
+			if err != nil {
 				return fmt.Errorf("write MPI hostfile line: %w", err)
 			}
 		}
diff --git a/runner/internal/runner/schemas/schemas.go b/runner/internal/runner/schemas/schemas.go
index 47706228cd..c9102d732d 100644
--- a/runner/internal/runner/schemas/schemas.go
+++ b/runner/internal/runner/schemas/schemas.go
@@ -95,6 +95,7 @@ type ClusterInfo struct {
 	JobIPs      []string `json:"job_ips"`
 	MasterJobIP string   `json:"master_job_ip"`
 	GPUSPerJob  int      `json:"gpus_per_job"`
+	GPUSPerNode []int    `json:"gpus_per_node"`
 }
 
 type SSHKey struct {
diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py
index df2d0b35d4..9ace4f474d 100644
--- a/src/dstack/_internal/cli/services/configurators/run.py
+++ b/src/dstack/_internal/cli/services/configurators/run.py
@@ -692,20 +692,25 @@ def register_commands_args(cls, parser: argparse.ArgumentParser):
             metavar="RUN_ARGS",
         )
 
-    def apply_commands_args(
-        self,
-        conf: ConfigurationWithCommandsParams,
-        args: argparse.Namespace,
-    ):
-        commands = conf.commands
+    def _interpolate_commands(self, commands: list[str], args: argparse.Namespace) -> None:
         run_args = shlex.join(args.run_args)
-        interpolator = VariablesInterpolator({"run": {"args": run_args}}, skip=["secrets"])
+        interpolator = VariablesInterpolator(
+            {"run": {"args": run_args}},
+            skip=["secrets", "groups"],
+        )
         try:
             for i, command in enumerate(commands):
                 commands[i] = interpolator.interpolate_or_error(command)
         except InterpolatorError as e:
             raise ConfigurationError(e.args[0])
 
+    def apply_commands_args(
+        self,
+        conf: ConfigurationWithCommandsParams,
+        args: argparse.Namespace,
+    ):
+        self._interpolate_commands(conf.commands, args)
+
 
 class TaskConfigurator(
     RunWithPortsConfiguratorMixin, RunWithCommandsConfiguratorMixin, BaseRunConfigurator
@@ -722,6 +727,9 @@ def apply_args(self, conf: TaskConfiguration, args: argparse.Namespace):
         super().apply_args(conf, args)
         self.apply_ports_args(conf, args)
         self.apply_commands_args(conf, args)
+        if conf.groups is not None:
+            for group in conf.groups:
+                self._interpolate_commands(group.commands, args)
 
 
 class DevEnvironmentConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigurator):
diff --git a/src/dstack/_internal/cli/utils/run.py b/src/dstack/_internal/cli/utils/run.py
index 6c27f2aa6f..46dfba5803 100644
--- a/src/dstack/_internal/cli/utils/run.py
+++ b/src/dstack/_internal/cli/utils/run.py
@@ -122,7 +122,7 @@ def th(s: str) -> str:
     props.add_row(th("User"), run_plan.user)
     configuration_type = run_spec.configuration.type
     if run_spec.configuration.type == "task":
-        configuration_type += f" (nodes={run_spec.configuration.nodes})"
+        configuration_type += f" (nodes={run_spec.configuration.nodes_num})"
     props.add_row(th("Type"), configuration_type)
     props.add_row(th("Resources"), pretty_req)
     props.add_row(th("Spot policy"), spot_policy)
diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py
index 3ae14d5d99..423de1f05d 100644
--- a/src/dstack/_internal/core/backends/slurm/compute.py
+++ b/src/dstack/_internal/core/backends/slurm/compute.py
@@ -164,6 +164,7 @@ def run_jobs(
             instance_offer=instance_offer,
             project_ssh_public_key=project_ssh_public_key,
             requirements=requirements,
+            node_count=len(job_configurations),
         )
 
     def terminate_instance(
@@ -186,6 +187,7 @@ def _run_slurm_job(
         instance_offer: InstanceOfferWithAvailability,
         project_ssh_public_key: str,
         requirements: Requirements,
+        node_count: Optional[int] = None,
     ) -> ComputeGroupProvisioningData:
         if job.job_spec.registry_auth is not None:
             self._skip_offer_cache.add(run, job, instance_offer)
@@ -209,7 +211,11 @@ def _run_slurm_job(
             assert run.run_spec.ssh_key_pub is not None
             authorized_keys = [project_ssh_public_key.strip(), run.run_spec.ssh_key_pub.strip()]
 
-            node_count = job.job_spec.jobs_per_replica
+            # Heterogeneous groups provision one shape at a time; Slurm allocation
+            # size must match that batch. Fall back to jobs_per_replica for
+            # run_job / homogeneous single-call paths.
+            if node_count is None:
+                node_count = job.job_spec.jobs_per_replica
             resources_spec = requirements.resources
             requested_resources = get_requested_resources_from_resources_spec(resources_spec)
 
diff --git a/src/dstack/_internal/core/compatibility/runs.py b/src/dstack/_internal/core/compatibility/runs.py
index 847e7be303..fa95dcd689 100644
--- a/src/dstack/_internal/core/compatibility/runs.py
+++ b/src/dstack/_internal/core/compatibility/runs.py
@@ -8,6 +8,7 @@
 )
 from dstack._internal.core.models.configurations import (
     ServiceConfiguration,
+    TaskConfiguration,
 )
 from dstack._internal.core.models.routers import SGLangServiceRouterConfig
 from dstack._internal.core.models.runs import (
@@ -98,6 +99,10 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType:
     if not run_spec.configuration.dstack:
         configuration_excludes["dstack"] = True
 
+    if isinstance(run_spec.configuration, TaskConfiguration):
+        if run_spec.configuration.groups is None:
+            configuration_excludes["groups"] = True
+
     if isinstance(run_spec.configuration, ServiceConfiguration):
         if run_spec.configuration.probes:
             probe_excludes: IncludeExcludeDictType = {}
@@ -160,6 +165,12 @@ def get_job_spec_excludes(job_specs: list[JobSpec]) -> IncludeExcludeDictType:
     spec_excludes: IncludeExcludeDictType = {}
     if all(s.replica_group == DEFAULT_REPLICA_GROUP_NAME for s in job_specs):
         spec_excludes["replica_group"] = True
+    if all(s.node_group_index == 0 for s in job_specs):
+        spec_excludes["node_group_index"] = True
+    if all(s.node_group_name == DEFAULT_REPLICA_GROUP_NAME for s in job_specs):
+        spec_excludes["node_group_name"] = True
+    if all(s.node_group_job_index == 0 for s in job_specs):
+        spec_excludes["node_group_job_index"] = True
 
     probe_excludes: IncludeExcludeDictType = {}
     spec_excludes["probes"] = {"__all__": probe_excludes}
diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py
index c05bd7f1db..04ea6a5385 100644
--- a/src/dstack/_internal/core/models/configurations.py
+++ b/src/dstack/_internal/core/models/configurations.py
@@ -708,6 +708,9 @@ def check_image_or_commands_present(self) -> Self:
         replicas = getattr(self, "replicas", None)
         if isinstance(replicas, list):
             return self
+        # If groups is set, skip validation - commands come from node groups
+        if getattr(self, "groups", None) is not None:
+            return self
 
         if not self.commands and not getattr(self, "image", None):
             raise ValueError("Either `commands` or `image` must be set")
@@ -798,8 +801,85 @@ def validate_dstack_and_inactivity_duration(self) -> Self:
         return self
 
 
+class NodeGroup(CoreModel):
+    name: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The name of the node group. If not provided, defaults to '0', '1', etc. "
+                "based on position."
+            )
+        ),
+    ] = None
+    nodes: Annotated[int, Field(description="The number of nodes in this group", ge=1)] = 1
+    resources: Annotated[
+        ResourcesSpec,
+        Field(description="The resources requirements for nodes in this group"),
+    ] = ResourcesSpec()
+    commands: Annotated[
+        CommandsList,
+        Field(description="The shell commands to run for nodes in this group"),
+    ] = []
+    ports: Annotated[
+        List[PortMappingOrShorthand],
+        Field(description="Port numbers/mapping to expose for nodes in this group"),
+    ] = []
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, v: Optional[str]) -> Optional[str]:
+        if v is not None:
+            if not is_valid_replica_group_name(v):
+                raise ValueError("Resource name should match regex '^[a-z0-9][a-z0-9-]{0,39}$'")
+        return v
+
+
 class TaskConfigurationParams(CoreModel):
-    nodes: Annotated[int, Field(description="Number of nodes", ge=1)] = 1
+    nodes: Annotated[
+        int,
+        Field(description="The number of nodes for homogeneous multi-node tasks", ge=1),
+    ] = 1
+    groups: Annotated[
+        Optional[List[NodeGroup]],
+        Field(
+            description=(
+                "A list of node groups for heterogeneous multi-node tasks. "
+                "Mutually exclusive with `nodes`."
+            ),
+        ),
+    ] = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_nodes_xor_groups(cls, data):
+        if not isinstance(data, dict):
+            return data
+        # Allow groups with default nodes: 1 (serialized configs always include it).
+        # Reject nodes: N (N != 1) together with groups.
+        if data.get("groups") is not None and "nodes" in data:
+            nodes = data.get("nodes")
+            if nodes is not None and nodes != 1:
+                raise ValueError("`nodes` and `groups` are mutually exclusive")
+        return data
+
+    @field_validator("groups")
+    @classmethod
+    def validate_groups(cls, v: Optional[List[NodeGroup]]) -> Optional[List[NodeGroup]]:
+        if v is None:
+            return v
+        if not v:
+            raise ValueError("`groups` cannot be an empty list")
+        for index, group in enumerate(v):
+            if group.name is None:
+                group.name = str(index)
+        counts = Counter(group.name for group in v)
+        duplicates = [name for name, count in counts.items() if count > 1]
+        if duplicates:
+            raise ValueError(
+                f"Duplicate node group names found: {duplicates}. "
+                "Each node group must have a unique name."
+            )
+        return v
 
 
 class TaskConfiguration(
@@ -811,6 +891,24 @@ class TaskConfiguration(
 ):
     type: Literal["task"] = "task"
 
+    @property
+    def node_groups(self) -> List[NodeGroup]:
+        if self.groups is not None:
+            return self.groups
+        return [
+            NodeGroup(
+                name=DEFAULT_REPLICA_GROUP_NAME,
+                nodes=self.nodes,
+                commands=self.commands,
+                resources=self.resources,
+                ports=self.ports,
+            )
+        ]
+
+    @property
+    def nodes_num(self) -> int:
+        return sum(group.nodes for group in self.node_groups)
+
 
 def _validate_replica_range(v: Range[int]) -> Range[int]:
     """Validate a Range[int] used for replica counts."""
diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py
index a292a928b8..c41c16084a 100644
--- a/src/dstack/_internal/core/models/runs.py
+++ b/src/dstack/_internal/core/models/runs.py
@@ -299,6 +299,18 @@ class JobSpec(CoreModel):
     service_port: Optional[int] = None
     """`service_port` is `None` for non-services and pre-0.19.19 services. See `get_service_port`."""
     probes: list[ProbeSpec] = []
+    node_group_index: int = 0
+    """`node_group_index` uses a default value for backward compatibility."""
+    node_group_name: str = DEFAULT_REPLICA_GROUP_NAME
+    """`node_group_name` uses a default value for backward compatibility."""
+    node_group_job_index: int = 0
+    """That node's index inside its group (0 .. group.nodes-1).
+    Example:
+      groups:
+        - nodes: 2   # jobs get node_group_job_index 0 and 1
+        - nodes: 1   # job gets node_group_job_index 0
+    Default for backward compatibility.
+    """
 
 
 class JobProvisioningData(CoreModel):
@@ -391,6 +403,12 @@ class ClusterInfo(CoreModel):
     job_ips: List[str]
     master_job_ip: str
     gpus_per_job: int
+    """GPU count on this node only."""
+    gpus_per_node: List[int] = []
+    """GPU count for each node in the run, in `job_ips` order.
+    Used for heterogeneous node groups where nodes can have different GPU
+    counts (e.g. `[2, 8]`). `0` means CPU-only. Empty for older servers.
+    """
 
 
 class Probe(CoreModel):
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py
index 720b0141ba..21d78da495 100644
--- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py
+++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py
@@ -121,6 +121,10 @@
 from dstack._internal.utils.common import get_current_datetime, get_or_error, run_async
 from dstack._internal.utils.interpolator import InterpolatorError
 from dstack._internal.utils.logging import get_logger
+from dstack._internal.utils.nodes_interpolator import (
+    find_groups_ip_refs,
+    interpolate_groups_ip_address,
+)
 
 logger = get_logger(__name__)
 
@@ -611,6 +615,28 @@ async def _prepare_startup_context(
         )
         return None
 
+    commands = context.job.job_spec.commands
+    if any(find_groups_ip_refs(c) for c in commands):
+        nodes_view = _build_nodes_ip_view(context.run.jobs, context.job.job_spec.replica_num)
+        try:
+            if not _referenced_ips_ready(commands, nodes_view):
+                logger.debug(
+                    "%s: waiting for referenced node group IPs",
+                    fmt(context.job_model),
+                )
+                return None
+            context.job.job_spec.commands = [
+                interpolate_groups_ip_address(c, nodes_view) for c in commands
+            ]
+        except InterpolatorError as e:
+            _terminate_job(
+                job_model=context.job_model,
+                job_update_map=result.job_update_map,
+                termination_reason=JobTerminationReason.TERMINATED_BY_SERVER,
+                termination_reason_message=f"Groups IP interpolation error: {e.args[0]}",
+            )
+            return None
+
     return _StartupContext(
         cluster_info=cluster_info,
         volumes=volumes,
@@ -1762,25 +1788,69 @@ def _reset_disconnected_at(job_model: JobModel, result: _ProcessResult) -> None:
         result.job_update_map["disconnected_at"] = None
 
 
+def _build_nodes_ip_view(jobs: list[Job], replica_num: int) -> list[list[str]]:
+    replica_jobs = [job for job in jobs if job.job_spec.replica_num == replica_num]
+    if not replica_jobs:
+        return []
+    max_group_index = max(job.job_spec.node_group_index for job in replica_jobs)
+    nodes: list[list[str]] = [[] for _ in range(max_group_index + 1)]
+    for job in replica_jobs:
+        group_index = job.job_spec.node_group_index
+        local_index = job.job_spec.node_group_job_index
+        while len(nodes[group_index]) <= local_index:
+            nodes[group_index].append("")
+        ip = ""
+        if job.job_submissions:
+            jpd = job.job_submissions[-1].job_provisioning_data
+            if jpd is not None:
+                ip = jpd.internal_ip or ""
+        nodes[group_index][local_index] = ip
+    return nodes
+
+
+def _referenced_ips_ready(commands: list[str], nodes_view: list[list[str]]) -> bool:
+    for command in commands:
+        for group_index, node_index in find_groups_ip_refs(command):
+            if group_index >= len(nodes_view) or node_index >= len(nodes_view[group_index]):
+                raise InterpolatorError(
+                    f"Invalid reference groups[{group_index}].nodes[{node_index}].IP_ADDRESS: "
+                    "out of range"
+                )
+            # Wait until every referenced slot has a non-empty internal IP.
+            if not nodes_view[group_index][node_index]:
+                return False
+    return True
+
+
 def _get_cluster_info(
     jobs: list[Job],
     replica_num: int,
     job_provisioning_data: JobProvisioningData,
     job_runtime_data: Optional[JobRuntimeData],
 ) -> ClusterInfo:
-    job_ips = []
-    for job in jobs:
-        if job.job_spec.replica_num == replica_num:
-            job_ips.append(
-                get_or_error(job.job_submissions[-1].job_provisioning_data).internal_ip or ""
-            )
+    job_ips: list[str] = []
+    gpus_per_node: list[int] = []
+    replica_jobs = sorted(
+        (job for job in jobs if job.job_spec.replica_num == replica_num),
+        key=lambda j: j.job_spec.job_num,
+    )
+    for job in replica_jobs:
+        submission = job.job_submissions[-1]
+        jpd = get_or_error(submission.job_provisioning_data)
+        job_ips.append(jpd.internal_ip or "")
+        jrd = submission.job_runtime_data
+        if jrd is not None and jrd.offer is not None:
+            gpus_per_node.append(len(jrd.offer.instance.resources.gpus))
+        else:
+            gpus_per_node.append(len(jpd.instance_type.resources.gpus))
     gpus_per_job = len(job_provisioning_data.instance_type.resources.gpus)
     if job_runtime_data is not None and job_runtime_data.offer is not None:
         gpus_per_job = len(job_runtime_data.offer.instance.resources.gpus)
     return ClusterInfo(
         job_ips=job_ips,
-        master_job_ip=job_ips[0],
+        master_job_ip=job_ips[0] if job_ips else "",
         gpus_per_job=gpus_per_job,
+        gpus_per_node=gpus_per_node,
     )
 
 
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py
index 74e1031c5f..e7cf969c2a 100644
--- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py
+++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py
@@ -49,6 +49,7 @@
     JobTerminationReason,
     Requirements,
     Run,
+    RunSpec,
 )
 from dstack._internal.core.models.volumes import Volume
 from dstack._internal.core.services.profiles import get_termination
@@ -777,25 +778,45 @@ async def _fetch_run_model_for_submitted_job(
 ) -> RunModel:
     """Fetch run model with only the relevant latest-submission jobs.
 
+    Loading jobs is separate from provisioning them. job_num=0 may load all
+    siblings for coordination, but still provisions only its own node group.
+
     Only a small subset is needed depending on the job type:
-        * Master multinode: all same-replica jobs (for cluster provisioning and releasing sibling waits).
-        * Non-master: master job + current job (for master provisioning data lookup).
-        * Master single-node: current job only (no siblings needed).
+        * Multinode master (job_num=0): all same-replica jobs.
+        * First job in a node group (not job 0): job 0 + jobs in its group
+          (same shape batch).
+        * Other multinode jobs: job 0 + current job.
+        * Single-node master: current job only.
 
     Only the latest submission per (replica_num, job_num) is loaded since historical
     submissions are never accessed in submitted job processing.
     """
+    job_spec = get_job_spec(job_model)
     is_master = job_model.job_num == 0
-    is_multinode = get_job_spec(job_model).jobs_per_replica > 1
+    is_multinode = job_spec.jobs_per_replica > 1
 
     job_num_filters: list = []
-    if is_master and not is_multinode:
-        # Master single-node: only current job needed.
-        job_num_filters.append(JobModel.job_num == 0)
-    elif not is_master:
-        # Non-master: master job (for provisioning data) + current job.
+    if not is_multinode:
+        if is_master:
+            # Single-node master: only current job needed.
+            job_num_filters.append(JobModel.job_num == 0)
+        else:
+            # Non-master single-node should not happen; keep master + current for safety.
+            job_num_filters.append(JobModel.job_num.in_([0, job_model.job_num]))
+    elif is_master:
+        # Multinode master: load all jobs (fleet setup, release waiting_master_job).
+        # Provisioning still batches only this job's node group (same shape).
+        pass
+    elif job_spec.node_group_job_index == 0:
+        # First job in a node group (not job 0): job 0 + this group's jobs.
+        run_spec = await _get_run_spec(session, job_model.run_id)
+        group_job_nums = _job_nums_for_node_group(
+            run_spec.configuration, job_spec.node_group_index
+        )
+        job_num_filters.append(JobModel.job_num.in_(sorted({0, *group_job_nums})))
+    else:
+        # Other multinode jobs: job 0 + current job.
         job_num_filters.append(JobModel.job_num.in_([0, job_model.job_num]))
-    # else: master multinode — no job_num filter, load all jobs in replica.
 
     latest_submissions_sq = (
         select(
@@ -839,6 +860,21 @@ async def _fetch_run_model_for_submitted_job(
     return res.unique().scalar_one()
 
 
+async def _get_run_spec(session: AsyncSession, run_id: uuid.UUID) -> RunSpec:
+    res = await session.execute(select(RunModel.run_spec).where(RunModel.id == run_id))
+    return RunSpec.model_validate_json(res.scalar_one())
+
+
+def _job_nums_for_node_group(configuration, group_index: int) -> list[int]:
+    assert configuration.type == "task"
+    job_num = 0
+    for index, group in enumerate(configuration.node_groups):
+        if index == group_index:
+            return list(range(job_num, job_num + group.nodes))
+        job_num += group.nodes
+    raise ValueError(f"node_group_index {group_index} out of range")
+
+
 def _get_job_models_for_jobs(
     job_models: list[JobModel],
     jobs: list[Job],
@@ -2142,15 +2178,56 @@ def _hint_pipelines_fetch(
         pipeline_hinter.hint_fetch(FleetModel.__name__)
 
 
+def _is_node_group_master(job: Job, replica_jobs: list[Job]) -> bool:
+    """True if `job` has the lowest job_num among loaded jobs in its node group.
+
+    `job` must be in `replica_jobs`.
+    """
+    group_index = job.job_spec.node_group_index
+    group_job_nums = [
+        j.job_spec.job_num for j in replica_jobs if j.job_spec.node_group_index == group_index
+    ]
+    return job.job_spec.job_num == min(group_job_nums)
+
+
+def _job_needs_provisioning(job: Job) -> bool:
+    if not job.job_submissions:
+        return True
+    return job.job_submissions[-1].job_provisioning_data is None
+
+
 def _select_jobs_to_provision(job: Job, replica_jobs: list[Job], job_model: JobModel) -> list[Job]:
-    jobs_to_provision = [job]
-    if is_multinode_job(job) and is_master_job(job) and job_model.waiting_master_job is not None:
-        jobs_to_provision = replica_jobs
-    return jobs_to_provision
+    """Select jobs to launch in this provision attempt.
+
+    Homogeneous multinode (`nodes: N` → one node group) still batches the whole
+    replica on the group master (rank 0).
+
+    Heterogeneous node groups batch only jobs that share `node_group_index`, so
+    ComputeGroup backends (`run_jobs`) receive a single-shape offer set.
+
+    Global `waiting_master_job` is unchanged: non-masters stay blocked until the
+    global master (job_num=0) finishes its provision attempt.
+    """
+    if not is_multinode_job(job):
+        return [job]
+    # Legacy rows without the master-wait protocol: provision one-by-one only.
+    if job_model.waiting_master_job is None:
+        return [job]
+    if not _is_node_group_master(job, replica_jobs):
+        return [job]
+
+    group_index = job.job_spec.node_group_index
+    group_jobs = [
+        j
+        for j in replica_jobs
+        if j.job_spec.node_group_index == group_index and _job_needs_provisioning(j)
+    ]
+    return group_jobs if group_jobs else [job]
 
 
 def _get_required_targeted_instance_offers(context: _SubmittedJobContext) -> int:
-    if is_multinode_job(context.job) and is_master_job(context.job):
+    # Node-group masters (including non-zero groups) may batch multiple jobs.
+    if is_multinode_job(context.job) and len(context.jobs_to_provision) > 1:
         return len(context.jobs_to_provision)
     return 1
 
@@ -2160,9 +2237,15 @@ def _release_replica_jobs_from_master_wait(
     replica_job_models: list[JobModel],
     jobs_to_provision: list[Job],
 ) -> None:
-    if len(jobs_to_provision) > 1:
-        logger.debug("%s: allow replica jobs to be provisioned one-by-one", fmt(job_model))
-        for replica_job_model in replica_job_models:
+    # Global master may only provision its own node group (len == 1). Still release
+    # waiting workers so other groups can provision on later ticks.
+    if job_model.job_num != 0:
+        return
+    if not any(m.waiting_master_job for m in replica_job_models):
+        return
+    logger.debug("%s: allow replica jobs to be provisioned one-by-one", fmt(job_model))
+    for replica_job_model in replica_job_models:
+        if replica_job_model.waiting_master_job:
             replica_job_model.waiting_master_job = False
 
 
diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py
index 5ec38790dd..9ae8b38b3d 100644
--- a/src/dstack/_internal/server/services/jobs/configurators/base.py
+++ b/src/dstack/_internal/server/services/jobs/configurators/base.py
@@ -3,6 +3,7 @@
 import sys
 import threading
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from pathlib import PurePosixPath
 from typing import Dict, List, Optional
 
@@ -26,6 +27,7 @@
     LEGACY_REPO_DIR,
     OPENAI_MODEL_PROBE_TIMEOUT,
     HTTPHeaderSpec,
+    NodeGroup,
     PortMapping,
     ProbeConfig,
     PythonVersion,
@@ -94,6 +96,13 @@ def get_default_image(nvcc: bool = False) -> str:
     return f"{settings.DSTACK_DOCKER_BASE_IMAGE}:{settings.DSTACK_DOCKER_BASE_IMAGE_VERSION}-{'devel' if nvcc else 'base'}-ubuntu{settings.DSTACK_DOCKER_BASE_IMAGE_UBUNTU_VERSION}"
 
 
+@dataclass(frozen=True)
+class NodeGroupJobContext:
+    group: NodeGroup
+    group_index: int
+    job_index: int
+
+
 class JobConfigurator(ABC):
     TYPE: RunConfigurationType
 
@@ -116,7 +125,7 @@ async def get_job_specs(self, replica_num: int) -> List[JobSpec]:
         return [job_spec]
 
     @abstractmethod
-    def _shell_commands(self) -> List[str]:
+    def _shell_commands(self, node_group: Optional[NodeGroup] = None) -> List[str]:
         pass
 
     @abstractmethod
@@ -135,7 +144,7 @@ def _reservation(self) -> Optional[str]:
         return self.run_spec.merged_profile.reservation
 
     @abstractmethod
-    def _ports(self) -> List[PortMapping]:
+    def _ports(self, node_group: Optional[NodeGroup] = None) -> List[PortMapping]:
         pass
 
     async def _get_image_config(self) -> ImageConfig:
@@ -165,15 +174,17 @@ async def _get_job_spec(
         replica_num: int,
         job_num: int,
         jobs_per_replica: int,
+        node_group_context: Optional[NodeGroupJobContext] = None,
     ) -> JobSpec:
+        node_group = node_group_context.group if node_group_context is not None else None
         job_spec = JobSpec(
             replica_num=replica_num,  # TODO(egor-s): add to env variables in the runner
             job_num=job_num,
             job_name=f"{self.run_spec.run_name}-{job_num}-{replica_num}",
             jobs_per_replica=jobs_per_replica,
             replica_group=self.replica_group_name or DEFAULT_REPLICA_GROUP_NAME,
-            app_specs=self._app_specs(),
-            commands=await self._commands(),
+            app_specs=self._app_specs(node_group),
+            commands=await self._commands(node_group),
             env=self._env(),
             home_dir=self._home_dir(),
             image_name=self._image_name(),
@@ -184,7 +195,7 @@ async def _get_job_spec(
             stop_duration=self._stop_duration(),
             utilization_policy=self._utilization_policy(),
             registry_auth=self._registry_auth(),
-            requirements=self._requirements(jobs_per_replica),
+            requirements=self._requirements(jobs_per_replica, node_group),
             retry=self._retry(),
             working_dir=self._working_dir(),
             volumes=self._volumes(job_num),
@@ -196,6 +207,17 @@ async def _get_job_spec(
             file_archives=self.run_spec.file_archives,
             service_port=self._service_port(),
             probes=self._probes(),
+            node_group_index=(
+                node_group_context.group_index if node_group_context is not None else 0
+            ),
+            node_group_name=(
+                node_group.name
+                if node_group is not None and node_group.name is not None
+                else DEFAULT_REPLICA_GROUP_NAME
+            ),
+            node_group_job_index=(
+                node_group_context.job_index if node_group_context is not None else 0
+            ),
         )
         return job_spec
 
@@ -210,12 +232,12 @@ def _shell(self) -> str:
             return "/bin/bash"
         return "/bin/sh"
 
-    async def _commands(self) -> List[str]:
+    async def _commands(self, node_group: Optional[NodeGroup] = None) -> List[str]:
         if self.run_spec.configuration.entrypoint is not None:  # docker-like format
             assert self.run_spec.configuration.type != "dev-environment"
             entrypoint = shlex.split(self.run_spec.configuration.entrypoint)
             commands = self.run_spec.configuration.commands
-        elif shell_commands := self._shell_commands():
+        elif shell_commands := self._shell_commands(node_group):
             entrypoint = [self._shell(), "-i", "-c"]
             dstack_image_commands = self._dstack_image_commands()
             commands = [_join_shell_commands(dstack_image_commands + shell_commands)]
@@ -265,9 +287,9 @@ def _dstack_image_commands(self) -> List[str]:
             f"eval $(echo '. $DSTACK_VENV_DIR/bin/activate' | sudo tee -a {DSTACK_PROFILE_PATH})",
         ]
 
-    def _app_specs(self) -> List[AppSpec]:
+    def _app_specs(self, node_group: Optional[NodeGroup] = None) -> List[AppSpec]:
         specs = []
-        for i, pm in enumerate(filter_reserved_ports(self._ports())):
+        for i, pm in enumerate(filter_reserved_ports(self._ports(node_group))):
             specs.append(
                 AppSpec(
                     port=pm.container_port,
@@ -335,13 +357,19 @@ def _utilization_policy(self) -> Optional[UtilizationPolicy]:
     def _registry_auth(self) -> Optional[RegistryAuth]:
         return self.run_spec.configuration.registry_auth
 
-    def _requirements(self, jobs_per_replica: int) -> Requirements:
+    def _requirements(
+        self,
+        jobs_per_replica: int,
+        node_group: Optional[NodeGroup] = None,
+    ) -> Requirements:
         resources = self.run_spec.configuration.resources
         if self.run_spec.configuration.type == "service":
             for group in self.run_spec.configuration.replica_groups:
                 if group.name == self.replica_group_name:
                     resources = group.resources
                     break
+        elif self.run_spec.configuration.type == "task" and node_group is not None:
+            resources = node_group.resources
         spot_policy = self._spot_policy()
         return Requirements(
             resources=resources,
diff --git a/src/dstack/_internal/server/services/jobs/configurators/dev.py b/src/dstack/_internal/server/services/jobs/configurators/dev.py
index e4ee0a2d56..39d77d63a2 100644
--- a/src/dstack/_internal/server/services/jobs/configurators/dev.py
+++ b/src/dstack/_internal/server/services/jobs/configurators/dev.py
@@ -1,7 +1,11 @@
 from typing import Dict, List, Optional
 
 from dstack._internal.core.errors import ServerClientError
-from dstack._internal.core.models.configurations import PortMapping, RunConfigurationType
+from dstack._internal.core.models.configurations import (
+    NodeGroup,
+    PortMapping,
+    RunConfigurationType,
+)
 from dstack._internal.core.models.profiles import SpotPolicy
 from dstack._internal.core.models.runs import RunSpec
 from dstack._internal.server.services.ides import get_ide
@@ -33,7 +37,7 @@ def __init__(
             self.ide = ide
         super().__init__(run_spec=run_spec, secrets=secrets, replica_group_name=replica_group_name)
 
-    def _shell_commands(self) -> List[str]:
+    def _shell_commands(self, node_group: Optional[NodeGroup] = None) -> List[str]:
         assert self.run_spec.configuration.type == "dev-environment"
 
         commands = []
@@ -65,6 +69,6 @@ def _default_max_duration(self) -> Optional[int]:
     def _spot_policy(self) -> SpotPolicy:
         return self.run_spec.merged_profile.spot_policy or SpotPolicy.ONDEMAND
 
-    def _ports(self) -> List[PortMapping]:
+    def _ports(self, node_group: Optional[NodeGroup] = None) -> List[PortMapping]:
         assert self.run_spec.configuration.type == "dev-environment"
         return self.run_spec.configuration.ports
diff --git a/src/dstack/_internal/server/services/jobs/configurators/service.py b/src/dstack/_internal/server/services/jobs/configurators/service.py
index 45bc4c8f72..9861e6fbf5 100644
--- a/src/dstack/_internal/server/services/jobs/configurators/service.py
+++ b/src/dstack/_internal/server/services/jobs/configurators/service.py
@@ -2,6 +2,7 @@
 
 from dstack._internal import settings
 from dstack._internal.core.models.configurations import (
+    NodeGroup,
     PortMapping,
     ReplicaGroup,
     RunConfigurationType,
@@ -24,7 +25,7 @@ def _current_replica_group(self) -> Optional[ReplicaGroup]:
                 return group
         return None
 
-    def _shell_commands(self) -> List[str]:
+    def _shell_commands(self, node_group: Optional[NodeGroup] = None) -> List[str]:
         assert self.run_spec.configuration.type == "service"
         group = self._current_replica_group()
         if group is not None:
@@ -124,5 +125,5 @@ def _reservation(self) -> Optional[str]:
             return group.reservation
         return super()._reservation()
 
-    def _ports(self) -> List[PortMapping]:
+    def _ports(self, node_group: Optional[NodeGroup] = None) -> List[PortMapping]:
         return []
diff --git a/src/dstack/_internal/server/services/jobs/configurators/task.py b/src/dstack/_internal/server/services/jobs/configurators/task.py
index 51c136dfe7..03bdf9e8cb 100644
--- a/src/dstack/_internal/server/services/jobs/configurators/task.py
+++ b/src/dstack/_internal/server/services/jobs/configurators/task.py
@@ -1,9 +1,16 @@
 from typing import List, Optional
 
-from dstack._internal.core.models.configurations import PortMapping, RunConfigurationType
+from dstack._internal.core.models.configurations import (
+    NodeGroup,
+    PortMapping,
+    RunConfigurationType,
+)
 from dstack._internal.core.models.profiles import SpotPolicy
 from dstack._internal.core.models.runs import JobSpec
-from dstack._internal.server.services.jobs.configurators.base import JobConfigurator
+from dstack._internal.server.services.jobs.configurators.base import (
+    JobConfigurator,
+    NodeGroupJobContext,
+)
 
 
 class TaskJobConfigurator(JobConfigurator):
@@ -11,18 +18,31 @@ class TaskJobConfigurator(JobConfigurator):
 
     async def get_job_specs(self, replica_num: int) -> List[JobSpec]:
         assert self.run_spec.configuration.type == "task"
+        groups = self.run_spec.configuration.node_groups
+        total = sum(group.nodes for group in groups)
+
         job_specs = []
-        for job_num in range(self.run_spec.configuration.nodes):
-            job_spec = await self._get_job_spec(
-                replica_num=replica_num,
-                job_num=job_num,
-                jobs_per_replica=self.run_spec.configuration.nodes,
-            )
-            job_specs.append(job_spec)
+        job_num = 0
+        for group_index, group in enumerate(groups):
+            for local_index in range(group.nodes):
+                job_spec = await self._get_job_spec(
+                    replica_num=replica_num,
+                    job_num=job_num,
+                    jobs_per_replica=total,
+                    node_group_context=NodeGroupJobContext(
+                        group=group,
+                        group_index=group_index,
+                        job_index=local_index,
+                    ),
+                )
+                job_specs.append(job_spec)
+                job_num += 1
         return job_specs
 
-    def _shell_commands(self) -> List[str]:
+    def _shell_commands(self, node_group: Optional[NodeGroup] = None) -> List[str]:
         assert self.run_spec.configuration.type == "task"
+        if node_group is not None and node_group.commands:
+            return node_group.commands
         return self.run_spec.configuration.commands
 
     def _default_single_branch(self) -> bool:
@@ -34,6 +54,8 @@ def _default_max_duration(self) -> Optional[int]:
     def _spot_policy(self) -> SpotPolicy:
         return self.run_spec.merged_profile.spot_policy or SpotPolicy.ONDEMAND
 
-    def _ports(self) -> List[PortMapping]:
+    def _ports(self, node_group: Optional[NodeGroup] = None) -> List[PortMapping]:
         assert self.run_spec.configuration.type == "task"
+        if node_group is not None and node_group.ports:
+            return node_group.ports
         return self.run_spec.configuration.ports
diff --git a/src/dstack/_internal/server/services/runs/__init__.py b/src/dstack/_internal/server/services/runs/__init__.py
index 02b72c981f..3448392ce9 100644
--- a/src/dstack/_internal/server/services/runs/__init__.py
+++ b/src/dstack/_internal/server/services/runs/__init__.py
@@ -1182,7 +1182,7 @@ async def _validate_run_volumes(
     # that won't be created immediately (e.g. range of replicas or nodes).
     nodes = 1
     if run_spec.configuration.type == "task":
-        nodes = run_spec.configuration.nodes
+        nodes = run_spec.configuration.nodes_num
     for job_num in range(nodes):
         volumes = await get_job_configured_volumes(
             session=session, project=project, run_spec=run_spec, job_num=job_num
diff --git a/src/dstack/_internal/server/services/runs/spec.py b/src/dstack/_internal/server/services/runs/spec.py
index 364f81769c..508b644ccf 100644
--- a/src/dstack/_internal/server/services/runs/spec.py
+++ b/src/dstack/_internal/server/services/runs/spec.py
@@ -256,7 +256,7 @@ def can_update_run_spec(current_run_spec: RunSpec, new_run_spec: RunSpec) -> boo
 def get_nodes_required_num(run_spec: RunSpec) -> int:
     nodes_required_num = 1
     if run_spec.configuration.type == "task":
-        nodes_required_num = run_spec.configuration.nodes
+        nodes_required_num = run_spec.configuration.nodes_num
     elif run_spec.configuration.type == "service":
         nodes_required_num = sum(
             group.count.min or 0 for group in run_spec.configuration.replica_groups
diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py
index 46b51a189e..079ec6ee15 100644
--- a/src/dstack/_internal/server/testing/common.py
+++ b/src/dstack/_internal/server/testing/common.py
@@ -447,10 +447,16 @@ async def create_job(
     if deployment_num is None:
         deployment_num = run.deployment_num
     run_spec = validate_json_extra_ignore(RunSpec, run.run_spec)
-    job_spec = (
-        await get_job_specs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=replica_num)
-    )[0]
-    job_spec.job_num = job_num
+    job_specs = await get_job_specs_from_run_spec(
+        run_spec=run_spec, secrets={}, replica_num=replica_num
+    )
+    if 0 <= job_num < len(job_specs):
+        job_spec = job_specs[job_num]
+    else:
+        job_spec = job_specs[0].model_copy(deep=True)
+        job_spec.job_num = job_num
+        job_spec.job_name = f"{run_spec.run_name}-{job_num}-{replica_num}"
+
     job = JobModel(
         project_id=run.project_id,
         fleet=fleet,
diff --git a/src/dstack/_internal/utils/interpolator.py b/src/dstack/_internal/utils/interpolator.py
index 9a4e44659b..641b71dafb 100644
--- a/src/dstack/_internal/utils/interpolator.py
+++ b/src/dstack/_internal/utils/interpolator.py
@@ -63,10 +63,15 @@ def interpolate(
                 raise InterpolatorError(f"No pattern closing: {s[opening:]}")
 
             name = s[opening + len(Pattern.opening) : closing].strip()
-            if not self.validate_name(name):
-                raise InterpolatorError(f"Illegal reference name: {name}")
-            if name.split(".")[0] in self.skip:
+            # Skip before validate_name so non-standard refs (e.g. groups[0].nodes[0].IP_ADDRESS)
+            # can be left for later interpolators. Invalid skipped names without brackets
+            # (e.g. secrets.pass-word) still raise.
+            root = name.split(".")[0]
+            skip_ns = root.split("[")[0]
+            if skip_ns in self.skip and ("[" in root or self.validate_name(name)):
                 tokens.append(s[opening : closing + len(Pattern.closing)])
+            elif not self.validate_name(name):
+                raise InterpolatorError(f"Illegal reference name: {name}")
             elif name in self.variables:
                 tokens.append(self.variables[name])
             else:
diff --git a/src/dstack/_internal/utils/nodes_interpolator.py b/src/dstack/_internal/utils/nodes_interpolator.py
new file mode 100644
index 0000000000..84512aa00a
--- /dev/null
+++ b/src/dstack/_internal/utils/nodes_interpolator.py
@@ -0,0 +1,24 @@
+import re
+
+from dstack._internal.utils.interpolator import InterpolatorError
+
+_GROUPS_IP_REF = re.compile(r"\$\{\{\s*groups\[(\d+)\]\.nodes\[(\d+)\]\.IP_ADDRESS\s*\}\}")
+
+
+def find_groups_ip_refs(s: str) -> list[tuple[int, int]]:
+    return [(int(m.group(1)), int(m.group(2))) for m in _GROUPS_IP_REF.finditer(s)]
+
+
+def interpolate_groups_ip_address(s: str, nodes: list[list[str]]) -> str:
+    def repl(m: re.Match) -> str:
+        gi, ni = int(m.group(1)), int(m.group(2))
+        if gi >= len(nodes) or ni >= len(nodes[gi]):
+            raise InterpolatorError(
+                f"Invalid reference groups[{gi}].nodes[{ni}].IP_ADDRESS: out of range"
+            )
+        ip = nodes[gi][ni]
+        if not ip:
+            raise InterpolatorError(f"IP not available for groups[{gi}].nodes[{ni}].IP_ADDRESS")
+        return ip
+
+    return _GROUPS_IP_REF.sub(repl, s)
diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py
index 98f6ac2b5b..3eb92a39a9 100644
--- a/src/tests/_internal/core/models/test_configurations.py
+++ b/src/tests/_internal/core/models/test_configurations.py
@@ -916,3 +916,120 @@ def test_ide_optional(self):
     def test_version_requires_ide(self):
         with pytest.raises(ValueError, match="`version` requires `ide` to be set"):
             DevEnvironmentConfigurationParams(version="1.80.0")
+
+
+class TestNodeGroups:
+    def test_parses_int_nodes(self):
+        parsed = parse_run_configuration({"type": "task", "nodes": 2, "commands": ["true"]})
+        assert parsed.type == "task"
+        assert parsed.nodes == 2
+        assert parsed.groups is None
+        assert parsed.nodes_num == 2
+        assert len(parsed.node_groups) == 1
+        assert parsed.node_groups[0].nodes == 2
+        assert parsed.node_groups[0].name == "0"
+
+    def test_parses_groups_and_defaults_names(self):
+        parsed = parse_run_configuration(
+            {
+                "type": "task",
+                "image": "debian",
+                "groups": [
+                    {"nodes": 2, "commands": ["echo head"]},
+                    {"name": "workers", "nodes": 1, "commands": ["echo worker"]},
+                ],
+            }
+        )
+        assert parsed.type == "task"
+        assert parsed.groups is not None
+        assert parsed.nodes_num == 3
+        assert parsed.node_groups[0].name == "0"
+        assert parsed.node_groups[1].name == "workers"
+        assert parsed.node_groups[0].commands == ["echo head"]
+        assert parsed.node_groups[1].commands == ["echo worker"]
+
+    def test_accepts_default_nodes_with_groups(self):
+        # Serialized TaskConfiguration always includes nodes=1 (the field default).
+        # The xor validator must allow that so model round-trips succeed.
+        parsed = parse_run_configuration(
+            {
+                "type": "task",
+                "image": "debian",
+                "nodes": 1,
+                "groups": [
+                    {"nodes": 2, "commands": ["echo head"]},
+                    {"name": "workers", "nodes": 1, "commands": ["echo worker"]},
+                ],
+            }
+        )
+        assert parsed.type == "task"
+        assert parsed.nodes == 1
+        assert parsed.groups is not None
+        assert parsed.nodes_num == 3
+
+    def test_groups_round_trip_via_model_dump(self):
+        parsed = parse_run_configuration(
+            {
+                "type": "task",
+                "image": "debian",
+                "groups": [
+                    {"nodes": 2, "commands": ["echo head"], "ports": [8000]},
+                    {"name": "workers", "nodes": 1, "commands": ["echo worker"]},
+                ],
+            }
+        )
+        assert parsed.type == "task"
+        dumped = parsed.model_dump(mode="json")
+        assert dumped["nodes"] == 1
+        assert dumped["groups"] is not None
+
+        reparsed = parse_run_configuration(dumped)
+        assert reparsed.type == "task"
+        assert reparsed.nodes == 1
+        assert reparsed.nodes_num == 3
+        assert [g.name for g in reparsed.node_groups] == ["0", "workers"]
+        assert reparsed.node_groups[0].commands == ["echo head"]
+        assert reparsed.node_groups[0].ports[0].container_port == 8000
+        assert reparsed.node_groups[1].commands == ["echo worker"]
+
+    def test_rejects_auto_name_collision_with_explicit_name(self):
+        # Unnamed group at index 1 becomes "1", colliding with an explicit name "1".
+        with pytest.raises(ConfigurationError, match="Duplicate node group names"):
+            parse_run_configuration(
+                {
+                    "type": "task",
+                    "image": "debian",
+                    "groups": [
+                        {"name": "1", "nodes": 1, "commands": ["true"]},
+                        {"nodes": 1, "commands": ["true"]},
+                    ],
+                }
+            )
+
+    def test_rejects_duplicate_group_names(self):
+        with pytest.raises(ConfigurationError, match="Duplicate node group names"):
+            parse_run_configuration(
+                {
+                    "type": "task",
+                    "image": "debian",
+                    "groups": [
+                        {"name": "head", "nodes": 1, "commands": ["true"]},
+                        {"name": "head", "nodes": 1, "commands": ["true"]},
+                    ],
+                }
+            )
+
+    def test_rejects_empty_groups(self):
+        with pytest.raises(ConfigurationError, match="cannot be an empty list"):
+            parse_run_configuration({"type": "task", "image": "debian", "groups": []})
+
+    def test_rejects_nodes_and_groups_together(self):
+        with pytest.raises(ConfigurationError, match="mutually exclusive"):
+            parse_run_configuration(
+                {
+                    "type": "task",
+                    "image": "debian",
+                    "nodes": 2,
+                    "groups": [{"nodes": 1, "commands": ["true"]}],
+                }
+            )
diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py b/src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py
new file mode 100644
index 0000000000..6366a274bc
--- /dev/null
+++ b/src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py
@@ -0,0 +1,115 @@
+from datetime import datetime, timezone
+from uuid import uuid4
+
+import pytest
+
+from dstack._internal.core.models.runs import Job, JobSpec, JobSubmission
+from dstack._internal.server.background.pipeline_tasks.jobs_running import (
+    _build_nodes_ip_view,
+    _get_cluster_info,
+    _referenced_ips_ready,
+)
+from dstack._internal.server.testing.common import get_job_provisioning_data
+from dstack._internal.utils.interpolator import InterpolatorError
+
+
+def _job(
+    *,
+    job_num: int,
+    node_group_index: int,
+    node_group_job_index: int,
+    internal_ip: str,
+    gpu_count: int,
+) -> Job:
+    return Job.model_construct(
+        job_spec=JobSpec.model_construct(
+            replica_num=0,
+            job_num=job_num,
+            node_group_index=node_group_index,
+            node_group_job_index=node_group_job_index,
+            commands=[],
+        ),
+        job_submissions=[
+            JobSubmission.model_construct(
+                id=uuid4(),
+                submitted_at=datetime.now(timezone.utc),
+                job_provisioning_data=get_job_provisioning_data(
+                    internal_ip=internal_ip,
+                    gpu_count=gpu_count,
+                ),
+                job_runtime_data=None,
+            )
+        ],
+    )
+
+
+class TestGetClusterInfo:
+    def test_fills_gpus_per_node(self):
+        jobs = [
+            _job(
+                job_num=0,
+                node_group_index=0,
+                node_group_job_index=0,
+                internal_ip="10.0.0.1",
+                gpu_count=8,
+            ),
+            _job(
+                job_num=1,
+                node_group_index=1,
+                node_group_job_index=0,
+                internal_ip="10.0.0.2",
+                gpu_count=4,
+            ),
+        ]
+        this_jpd = get_job_provisioning_data(internal_ip="10.0.0.1", gpu_count=8)
+        info = _get_cluster_info(
+            jobs=jobs,
+            replica_num=0,
+            job_provisioning_data=this_jpd,
+            job_runtime_data=None,
+        )
+        assert info.job_ips == ["10.0.0.1", "10.0.0.2"]
+        assert info.master_job_ip == "10.0.0.1"
+        assert info.gpus_per_job == 8
+        assert info.gpus_per_node == [8, 4]
+
+
+class TestNodesIpView:
+    def test_builds_group_view(self):
+        jobs = [
+            _job(
+                job_num=0,
+                node_group_index=0,
+                node_group_job_index=0,
+                internal_ip="10.0.0.1",
+                gpu_count=1,
+            ),
+            _job(
+                job_num=1,
+                node_group_index=0,
+                node_group_job_index=1,
+                internal_ip="10.0.0.2",
+                gpu_count=1,
+            ),
+            _job(
+                job_num=2,
+                node_group_index=1,
+                node_group_job_index=0,
+                internal_ip="10.0.0.3",
+                gpu_count=1,
+            ),
+        ]
+        assert _build_nodes_ip_view(jobs, replica_num=0) == [
+            ["10.0.0.1", "10.0.0.2"],
+            ["10.0.0.3"],
+        ]
+
+    def test_referenced_ips_ready(self):
+        nodes_view = [["10.0.0.1"], [""]]
+        assert _referenced_ips_ready(["echo ${{ groups[0].nodes[0].IP_ADDRESS }}"], nodes_view)
+        assert not _referenced_ips_ready(["echo ${{ groups[1].nodes[0].IP_ADDRESS }}"], nodes_view)
+
+    def test_referenced_ips_out_of_range(self):
+        nodes_view = [["10.0.0.1"]]
+        with pytest.raises(InterpolatorError, match="out of range"):
+            _referenced_ips_ready(["echo ${{ groups[1].nodes[0].IP_ADDRESS }}"], nodes_view)
diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py
index ee34ec5ad5..491868327f 100644
--- a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py
+++ b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py
@@ -17,7 +17,11 @@
     RegistryAuth,
     validate_json_extra_ignore,
 )
-from dstack._internal.core.models.configurations import ServiceConfiguration, TaskConfiguration
+from dstack._internal.core.models.configurations import (
+    NodeGroup,
+    ServiceConfiguration,
+    TaskConfiguration,
+)
 from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.fleets import FleetNodesSpec, InstanceGroupPlacement
 from dstack._internal.core.models.instances import InstanceStatus
@@ -2692,13 +2696,15 @@ async def test_single_node_master_loads_only_current_job(self, test_db, session:
         assert not context.multinode
         assert context.jobs_to_provision == [context.job]
 
-    async def test_non_master_loads_master_and_current_job(self, test_db, session: AsyncSession):
-        """Non-master: run_model.jobs should contain master job + current job (latest submissions)."""
+    async def test_non_master_multinode_loads_master_and_current_job(
+        self, test_db, session: AsyncSession
+    ):
+        """Homogeneous multinode workers: run_model.jobs should contain job 0 + current."""
         project = await create_project(session=session)
         user = await create_user(session=session)
         repo = await create_repo(session=session, project_id=project.id)
         fleet = await create_fleet(session=session, project=project)
-        configuration = TaskConfiguration(image="debian", nodes=2)
+        configuration = TaskConfiguration(image="debian", nodes=3)
         run_spec = get_run_spec(run_name="run", repo_id=repo.name, configuration=configuration)
         run = await create_run(
             session=session,
@@ -2718,19 +2724,25 @@ async def test_non_master_loads_master_and_current_job(self, test_db, session: A
             job_provisioning_data=get_job_provisioning_data(),
             waiting_master_job=False,
         )
-        worker_job = await create_job(
+        worker_job_1 = await create_job(
             session=session,
             run=run,
             job_num=1,
             status=JobStatus.SUBMITTED,
             waiting_master_job=False,
         )
+        await create_job(
+            session=session,
+            run=run,
+            job_num=2,
+            status=JobStatus.SUBMITTED,
+            waiting_master_job=False,
+        )
         await session.commit()
 
-        context = await _load_submitted_job_context(session=session, job_model=worker_job)
-        # Only master (job_num=0) and current job (job_num=1) should be loaded.
+        context = await _load_submitted_job_context(session=session, job_model=worker_job_1)
         loaded_job_ids = {jm.id for jm in context.run_model.jobs}
-        assert loaded_job_ids == {master_job.id, worker_job.id}
+        assert loaded_job_ids == {master_job.id, worker_job_1.id}
         assert context.jobs_to_provision == [context.job]
 
     async def test_multinode_master_loads_all_replica_jobs(self, test_db, session: AsyncSession):
@@ -2774,6 +2786,87 @@ async def test_multinode_master_loads_all_replica_jobs(self, test_db, session: A
         assert len(context.jobs_to_provision) == 2
         assert len(context.replica_job_model_ids) == 2
 
+    async def test_node_group_master_provisions_only_its_group(
+        self, test_db, session: AsyncSession
+    ):
+        """Heterogeneous node groups: each group master batches only its own group."""
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        repo = await create_repo(session=session, project_id=project.id)
+        fleet = await create_fleet(session=session, project=project)
+        configuration = TaskConfiguration(
+            image="debian",
+            groups=[
+                NodeGroup(
+                    name="small",
+                    nodes=1,
+                    resources=ResourcesSpec(
+                        cpu=Range[int](min=2), memory=Range[Memory](min=Memory(4))
+                    ),
+                    commands=["echo small"],
+                ),
+                NodeGroup(
+                    name="large",
+                    nodes=2,
+                    resources=ResourcesSpec(
+                        cpu=Range[int](min=4), memory=Range[Memory](min=Memory(8))
+                    ),
+                    commands=["echo large"],
+                ),
+            ],
+        )
+        run_spec = get_run_spec(run_name="run", repo_id=repo.name, configuration=configuration)
+        run = await create_run(
+            session=session,
+            run_name="run",
+            project=project,
+            repo=repo,
+            user=user,
+            run_spec=run_spec,
+            fleet=fleet,
+        )
+        small_job = await create_job(
+            session=session,
+            run=run,
+            job_num=0,
+            status=JobStatus.SUBMITTED,
+            waiting_master_job=False,
+        )
+        large_job_0 = await create_job(
+            session=session,
+            run=run,
+            job_num=1,
+            status=JobStatus.SUBMITTED,
+            waiting_master_job=False,
+        )
+        large_job_1 = await create_job(
+            session=session,
+            run=run,
+            job_num=2,
+            status=JobStatus.SUBMITTED,
+            waiting_master_job=False,
+        )
+        await session.commit()
+
+        small_context = await _load_submitted_job_context(session=session, job_model=small_job)
+        # Multinode master: load all replica jobs.
+        assert {jm.job_num for jm in small_context.run_model.jobs} == {0, 1, 2}
+        assert [j.job_spec.job_num for j in small_context.jobs_to_provision] == [0]
+        assert small_context.jobs_to_provision[0].job_spec.node_group_name == "small"
+
+        large_context = await _load_submitted_job_context(session=session, job_model=large_job_0)
+        # Heterogeneous group master: job 0 + its node group.
+        assert {jm.job_num for jm in large_context.run_model.jobs} == {0, 1, 2}
+        assert sorted(j.job_spec.job_num for j in large_context.jobs_to_provision) == [1, 2]
+        assert {j.job_spec.node_group_name for j in large_context.jobs_to_provision} == {"large"}
+
+        large_worker_context = await _load_submitted_job_context(
+            session=session, job_model=large_job_1
+        )
+        # Non-master job in the group: job 0 + current.
+        assert {jm.job_num for jm in large_worker_context.run_model.jobs} == {0, 2}
+        assert [j.job_spec.job_num for j in large_worker_context.jobs_to_provision] == [2]
+
     async def test_loads_only_latest_submission(self, test_db, session: AsyncSession):
         """Only the latest submission per (replica_num, job_num) should be loaded, not historical ones."""
         project = await create_project(session=session)
diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py
index 46d530fff1..9319baba9f 100644
--- a/src/tests/_internal/server/routers/test_runs.py
+++ b/src/tests/_internal/server/routers/test_runs.py
@@ -321,6 +321,9 @@ def get_dev_env_run_plan_dict(
                     "file_archives": [],
                     "service_port": None,
                     "probes": [],
+                    "node_group_index": 0,
+                    "node_group_name": "0",
+                    "node_group_job_index": 0,
                 },
                 "offers": [json.loads(o.model_dump_json()) for o in offers],
                 "total_offers": total_offers,
@@ -568,6 +571,9 @@ def get_dev_env_run_dict(
                     "file_archives": [],
                     "service_port": None,
                     "probes": [],
+                    "node_group_index": 0,
+                    "node_group_name": "0",
+                    "node_group_job_index": 0,
                 },
                 "job_submissions": [
                     {
diff --git a/src/tests/_internal/server/services/jobs/configurators/test_task.py b/src/tests/_internal/server/services/jobs/configurators/test_task.py
index 54b8dd666d..383cd411e2 100644
--- a/src/tests/_internal/server/services/jobs/configurators/test_task.py
+++ b/src/tests/_internal/server/services/jobs/configurators/test_task.py
@@ -3,7 +3,8 @@
 
 import pytest
 
-from dstack._internal.core.models.configurations import TaskConfiguration
+from dstack._internal.core.models.configurations import NodeGroup, TaskConfiguration
+from dstack._internal.core.models.resources import GPUSpec, ResourcesSpec
 from dstack._internal.core.models.runs import JobSSHKey
 from dstack._internal.server.services.docker import ImageConfig
 from dstack._internal.server.services.jobs.configurators.task import TaskJobConfigurator
@@ -37,6 +38,65 @@ async def test_multi_node(self):
         assert job_specs[1].ssh_key == JobSSHKey(private="private1", public="public1")
 
 
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("image_config_mock")
+class TestNodeGroups:
+    async def test_assigns_contiguous_ranks_and_metadata(self):
+        configuration = TaskConfiguration(
+            image="debian",
+            groups=[
+                NodeGroup(name="head", nodes=2, commands=["echo head"]),
+                NodeGroup(name="workers", nodes=2, commands=["echo worker"]),
+            ],
+        )
+        run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration)
+        configurator = TaskJobConfigurator(run_spec)
+
+        job_specs = await configurator.get_job_specs(replica_num=0)
+
+        assert len(job_specs) == 4
+        assert [j.job_num for j in job_specs] == [0, 1, 2, 3]
+        assert [j.jobs_per_replica for j in job_specs] == [4, 4, 4, 4]
+        assert [j.node_group_name for j in job_specs] == [
+            "head",
+            "head",
+            "workers",
+            "workers",
+        ]
+        assert [j.node_group_index for j in job_specs] == [0, 0, 1, 1]
+        assert [j.node_group_job_index for j in job_specs] == [0, 1, 0, 1]
+
+    async def test_uses_per_group_commands_and_resources(self):
+        configuration = TaskConfiguration(
+            image="debian",
+            groups=[
+                NodeGroup(
+                    name="head",
+                    nodes=1,
+                    commands=["echo head"],
+                    resources=ResourcesSpec(gpu=GPUSpec(name=["H100"], count=1)),
+                ),
+                NodeGroup(
+                    name="workers",
+                    nodes=1,
+                    commands=["echo worker"],
+                    resources=ResourcesSpec(gpu=GPUSpec(name=["A100"], count=2)),
+                ),
+            ],
+        )
+        run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration)
+        configurator = TaskJobConfigurator(run_spec)
+
+        job_specs = await configurator.get_job_specs(replica_num=0)
+
+        assert "echo head" in job_specs[0].commands[-1]
+        assert "echo worker" in job_specs[1].commands[-1]
+        assert job_specs[0].requirements.resources.gpu.name == ["H100"]
+        assert job_specs[0].requirements.resources.gpu.count.min == 1
+        assert job_specs[1].requirements.resources.gpu.name == ["A100"]
+        assert job_specs[1].requirements.resources.gpu.count.min == 2
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("image_config_mock")
 class TestServerAccess:
diff --git a/src/tests/_internal/utils/test_interpolator.py b/src/tests/_internal/utils/test_interpolator.py
index 50c3845832..2acc2eaefc 100644
--- a/src/tests/_internal/utils/test_interpolator.py
+++ b/src/tests/_internal/utils/test_interpolator.py
@@ -49,3 +49,8 @@ def test_illegal_name(self):
             get_interpolator().interpolate("${{ secrets.password.hash }}")
         with pytest.raises(InterpolatorError):
             get_interpolator().interpolate("${{ secrets.007 }}")
+
+    def test_skips_groups_refs(self):
+        s = "ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379"
+        interpolator = VariablesInterpolator({"run": {"args": "x"}}, skip=["groups"])
+        assert interpolator.interpolate(s) == s
diff --git a/src/tests/_internal/utils/test_nodes_interpolator.py b/src/tests/_internal/utils/test_nodes_interpolator.py
new file mode 100644
index 0000000000..84ed9cdc85
--- /dev/null
+++ b/src/tests/_internal/utils/test_nodes_interpolator.py
@@ -0,0 +1,40 @@
+import pytest
+
+from dstack._internal.utils.interpolator import InterpolatorError
+from dstack._internal.utils.nodes_interpolator import (
+    find_groups_ip_refs,
+    interpolate_groups_ip_address,
+)
+
+
+class TestFindGroupsIpRefs:
+    def test_finds_refs(self):
+        s = "ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379"
+        assert find_groups_ip_refs(s) == [(0, 0)]
+
+    def test_finds_multiple_refs(self):
+        s = "${{ groups[0].nodes[1].IP_ADDRESS }} ${{groups[2].nodes[0].IP_ADDRESS}}"
+        assert find_groups_ip_refs(s) == [(0, 1), (2, 0)]
+
+    def test_no_refs(self):
+        assert find_groups_ip_refs("echo hello") == []
+
+
+class TestInterpolateGroupsIpAddress:
+    def test_replaces_ip(self):
+        s = "ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379"
+        result = interpolate_groups_ip_address(s, [["10.0.0.1", "10.0.0.2"], ["10.0.0.3"]])
+        assert result == "ray start --address=10.0.0.1:6379"
+
+    def test_replaces_nested_node(self):
+        s = "${{ groups[1].nodes[0].IP_ADDRESS }}"
+        result = interpolate_groups_ip_address(s, [["10.0.0.1"], ["10.0.0.2"]])
+        assert result == "10.0.0.2"
+
+    def test_raises_when_ip_missing(self):
+        with pytest.raises(InterpolatorError, match="IP not available"):
+            interpolate_groups_ip_address("${{ groups[0].nodes[0].IP_ADDRESS }}", [[""]])
+
+    def test_raises_when_out_of_range(self):
+        with pytest.raises(InterpolatorError, match="out of range"):
+            interpolate_groups_ip_address("${{ groups[1].nodes[0].IP_ADDRESS }}", [["10.0.0.1"]])