Inefficient GPU allocation for multi-replica service (especially small models)

[JenniferWang]

Summary

The current num_hosts parameter in ProcessConfig/ServiceConfig forces host-level allocation, preventing efficient GPU packing when multiple replicas could share the same host. e.g. a model that works best for TP=4 (vLLM).

Current Behavior

service = await Generator.options(
    procs=4,           # 4 GPU processes per replica
    with_gpus=True,
    num_hosts=1,       # Each replica gets 1 dedicated host
    num_replicas=2     # Total: 2 hosts needed
).as_service()

Desired Behavior

Allow efficient GPU-level packing:

service = await Generator.options(
    procs=4,           # 4 GPU processes per replica
    with_gpus=True,
    num_replicas=2     # Provisioner packs onto available GPUs
).as_service()

Proposal

The current resource allocating semantic is a bit awkward to express these three specifications at the same time

1. HostMesh
2. ProcMesh (on which HostMesh)
3. ServiceProcMesh (on which HostMesh)

So the proposal is to change the current config semantic to be something along this line

device_per_host: 8 
host_meshs: # the following adds up to 7 hosts
   policy: 
      num_hosts: 4 
   ref_model:
      num_hosts: 1
   trainer: 
      num_hosts: 2

services:
   policy:
      procs: 4
      num_replicas: 8  # maximize the hosts
      with_gpus: true
      host_mesh: policy
   ref_model:
      procs: 4
      num_replicas: 1
      with_gpus: true
      host_mesh: ref_model
   reward_actor:
      procs: 4
      num_replicas: 1
      with_gpus: false
      host_mesh: default

actors:
   dataset:
      procs: 1
      with_gpus: false
      host_mesh: policy
   trainer:
      procs: 2
      with_gpus: true
      mesh_name: trainer
   replay_buffer:
      procs: 1
      with_gpus: false
      mesh_name: trainer
   compute_advantages:
      procs: 1
      with_gpus: false
      mesh_name: default

[allenwang28]

This is a good idea, I had a similar idea! As we're doing this though, I think we might want to also support colocation, i.e. a scenario where we have 2 hosts A and B, where trainer is scheduled on GPUs 0-3 on host A and generators are scheduled on the remaining GPUs.

How could we express that?

[JenniferWang]

This is a good idea, I had a similar idea! As we're doing this though, I think we might want to also support colocation, i.e. a scenario where we have 2 hosts A and B, where trainer is scheduled on GPUs 0-3 on host A and generators are scheduled on the remaining GPUs.

How could we express that?

Ooo, i didn't consider that, one option....ugly though

host_meshes:
   policy: 
     num_hosts: 1 
   shared:
     num_hosts: 1 

services:
   policy:
      procs: 4
      num_replicas: 3  # 1 replica on 'shared' and the rest on 'policy'
      with_gpus: true
      host_meshes: [policy, shared] 

Another thing i forgot to mention is -- we shouldn't need to expose non-gpu ProcMesh in the yaml config files (e.g. dataset, replay buffer etc).

[allenwang28]

here's a summary that Claude generated from what I was thinking - basically using slice semantics that Monarch provides

Proposal

Introduce HostAllocation (or RemoteHosts?) as a Forge abstraction that defines and creates Monarch HostMesh objects. It represents a remote host allocation with shape (num_hosts, gpus_per_host).

Examples

Case 1: Separate Allocation Per Service

# Allocate dedicated hosts for trainer
trainer_alloc = HostAllocation(num_hosts=1, gpus_per_host=8)

trainer = Trainer.options(
    pool=trainer_alloc,
    hosts=1,
    gpus=8
).as_actor()

# Allocate dedicated hosts for policy
policy_alloc = HostAllocation(num_hosts=3, gpus_per_host=8)

policy = Policy.options(
    pool=policy_alloc,
    hosts=1,
    gpus=2,
    replicas=12  # 3 hosts × (8 GPUs / 2 GPUs per replica)
).as_service()

Case 2: Auto-deduce Replicas from Allocation

When replicas is not specified, automatically fill the allocation

# Allocate N hosts with 8 GPUs each
allocation = HostAllocation(num_hosts=4, gpus_per_host=8)  # 32 total GPUs

# M=2 GPUs per replica
# Replicas auto-deduced: 4 hosts × (8 GPUs / 2) = 16 replicas
policy = Policy.options(
    pool=allocation,
    hosts=1,
    gpus=2
    # replicas automatically set to 16
).as_service()

# M=4 GPUs per replica  
# Replicas auto-deduced: 4 hosts × (8 GPUs / 4) = 8 replicas
policy = Policy.options(
    pool=allocation,
    hosts=1,
    gpus=4
    # replicas automatically set to 8
).as_service()

# M=8 GPUs per replica
# Replicas auto-deduced: 4 hosts × (8 GPUs / 8) = 4 replicas
policy = Policy.options(
    pool=allocation,
    hosts=1,
    gpus=8
    # replicas automatically set to 4
).as_service()

• gpus_per_replica must evenly divide gpus_per_host (i.e., gpus_per_host % gpus_per_replica == 0)
• If replicas is explicitly provided, validate it doesn't exceed allocation capacity

Case 3: Colocation with GPU-Level Slicing

Slice allocations to colocate services on the same physical hosts using different GPU ranges:

# 2 hosts (A and B), 8 GPUs each
allocation = HostAllocation(num_hosts=2, gpus_per_host=8)

# Trainer on GPUs 0-3 of host A
trainer_alloc = allocation.slice(hosts=0, gpus=slice(0, 4))  # (1, 4) allocation

trainer = Trainer.options(
    pool=trainer_alloc,  # Host A, GPUs 0-3
    hosts=1,
    gpus=4
).as_actor()

# Generators on GPUs 4-7 of host A  
generator_alloc = allocation.slice(hosts=0, gpus=slice(4, 8))  # (1, 4) allocation

generator = Generator.options(
    pool=generator_alloc,  # Host A, GPUs 4-7
    hosts=1,
    gpus=2
    # replicas auto-deduced: 1 host × (4 GPUs / 2) = 2 replicas
).as_service()

# Use host B for other services
host_b_alloc = allocation.slice(hosts=1)  # (1, 8) allocation

Across all hosts

allocation = HostAllocation(num_hosts=2, gpus_per_host=8)

# Trainer on GPUs 0-3 of ALL hosts
trainer_alloc = allocation.slice(hosts=slice(None), gpus=slice(0, 4))  # (2, 4) allocation

trainer = Trainer.options(
    pool=trainer_alloc,
    hosts=2,
    gpus=4  # 4 GPUs per host = 8 total
).as_actor()

# Generators on GPUs 4-7 of ALL hosts
generator_alloc = allocation.slice(hosts=slice(None), gpus=slice(4, 8))  # (2, 4) allocation

generator = Generator.options(
    pool=generator_alloc,
    hosts=1,
    gpus=2
    # replicas auto-deduced: 2 hosts × (4 GPUs / 2) = 4 replicas
).as_service()

What if we wanted to use torch_memory_saver? then we can re-use the same allocation:

# Single host with 8 GPUs
allocation = HostAllocation(num_hosts=1, gpus_per_host=8)

# Policy service uses all 8 GPUs
policy = Policy.options(
    pool=allocation,  # Uses all 8 GPUs
    hosts=1,
    gpus=8
).as_service()

# Reward service ALSO wants all 8 GPUs - same allocation!
reward = RewardModel.options(
    pool=allocation,  # Re-uses same allocation
    hosts=1,
    gpus=8
).as_service()

# Both services claim the same 8 GPUs, but coordinate via sleep/resume
async def training_loop():
    while True:
        # Phase 1: Generate rollouts with policy
        await reward.sleep()  # Reward releases GPU memory
        await policy.resume()  # Policy claims GPUs
        
        episodes = []
        for _ in range(batch_size):
            episode = await policy.generate.route(prompt)
            episodes.append(episode)
        
        # Phase 2: Score with reward model
        await policy.sleep()  # Policy releases GPU memory
        await reward.resume()  # Reward claims GPUs
        
        scores = await reward.score.fanout(episodes)
        
        # Phase 3: Train (both asleep, trainer uses GPUs)
        await reward.sleep()
        # Now trainer could use the GPUs...

Pros/Cons of HostAllocation Approach

Pros

1. Explicit and predictable: You know exactly what resources are allocated upfront
2. Flexible sharing: Supports both spatial partitioning (slicing) and temporal multiplexing (sleep/resume)
3. Monarch-aligned: Clean separation between resource definition (Forge) and mesh primitive (Monarch)
4. Composable: Can slice multiple times, reason about resource topology

Cons

1. Verbose for simple cases: Requires allocation boilerplate even for basic scenarios

   allocation = HostAllocation(num_hosts=1, gpus_per_host=8)
   policy = Policy.options(pool=allocation, hosts=1, gpus=8).as_service()

2. Manual coordination required: Sleep/resume is entirely user-managed - easy to mess up

   await policy.sleep()
   await reward.resume()  # Forget this? OOM crash

3. No over-allocation detection: Nothing prevents conflicting allocations - errors only at runtime

   # Both claim same 8 GPUs - no error until they both try to use them
   policy = Policy.options(pool=allocation, hosts=1, gpus=8).as_service()
   reward = Reward.options(pool=allocation, hosts=1, gpus=8).as_service()

4. Static allocation: Size fixed at creation - doesn't play well with autoscaling

   allocation = HostAllocation(num_hosts=4, gpus_per_host=8)
   # What if service wants to scale beyond 32 GPUs?

5. Boilerplate for common patterns: Temporal multiplexing requires explicit coordination everywhere

The Trade-off

Explicitness (Monarch philosophy) vs Ergonomics (RL common patterns)

Possible hybrid: explicit HostAllocation when needed, implicit allocation for simple cases.