[RFC]feat: Multi-policy Disaggregate Trainer

### RFC: Multi-policy (models) Fully Async Disaggregate Trainer

**Note**: This feature might be added to Verl instead of Uni-agent. The RFC is temporarily proposed at Uni-agent for convenience. It might be moved to Verl later.

## 1. Proposal

This RFC proposes a `MultiModelDisaggregateTrainer` to support multi-model co-training in Uni-Agent / VERL.

The core idea is simple:

```text
One model = one fully async disaggregate runtime.
MultiModelDisaggregateTrainer = launcher + lifecycle manager + training data consumer.
Gateway = request router.
```

Each model reuses the existing VERL fully async disaggregate trainer as much as possible. We do not want to build a new trainer from scratch. Instead, we wrap each single-model async trainer into a `ModelAsyncRuntime`, then let the top-level `MultiModelDisaggregateTrainer` start and manage multiple runtimes.

For model serving, we choose **Gateway-side routing**:

```text
Agent request
  -> Gateway
  -> AgentID -> ModelID
  -> corresponding LLM Server Client
```

The Multi-model Trainer should **not** expose one unified LLM Server Client and route requests internally. The trainer should start the model runtimes, collect their LLM Server Client handles, and pass these handles to Gateway. Gateway then directly calls the right model client.

---

## 2. Motivation

### 2.1 Improve rollout efficiency

In colocate mode, training and rollout are tightly coupled. If we have multiple models, their training speed can be different. For example, model A may finish training quickly, while model B is still running optimizer steps. In colocate mode, even if model A is ready, the whole agent task may still be unable to rollout because another model is still training.

With disaggregate mode, training and rollout are separated. Except for the short weight sync window, rollout servers stay available most of the time. This is more efficient for multi-agent workloads.

```text
Colocate:
  train all models -> wait slowest model -> rollout

Disaggregate:
  each model trains independently
  rollout stays available unless that model is syncing weights
```

### 2.2 Reduce long-tail effect

Colocate mode already has a long-tail issue in single-model training. Some rollout samples are slow, and the whole batch waits for them.

In multi-agent training, this becomes worse because long-tail effects can stack across agents and models. One slow model call or one slow agent trajectory may block the whole pipeline.

Async + disaggregate reduces this issue because rollout and training become streaming-style. Different models can progress independently, and we do not need to wait for every model at every step.

### 2.3 Better fit for dynamic agent workflows

In dynamic agent workflows, each model may be called a different number of times.

For example:

```text
planner model: called once
tool-use model: called many times
verifier model: called only for hard samples
```

Their batch collection speed will be different. If everything is colocated, models will frequently wait for each other. This is not ideal.

Disaggregate mode is more natural here. Each model has its own rollout service, training queue, trainer, and sync schedule. High-frequency models can train more often, while low-frequency models can train with smaller batches or lower update frequency.

---

## 3. Architecture Decision: Gateway Routes Directly to Model Clients

There are two possible designs.

### Option A: Gateway directly calls each model’s LLM Server Client

```text
MultiModelTrainer
  -> starts ModelRuntime[A]
  -> starts ModelRuntime[B]
  -> gets LLMServerClient[A], LLMServerClient[B]
  -> registers them to Gateway

Gateway
  -> receives AgentID
  -> maps AgentID to ModelID
  -> calls LLMServerClient[ModelID]
```

### Option B: Gateway calls one unified client from MultiModelTrainer

```text
Gateway
  -> calls UnifiedMultiModelClient
  -> MultiModelTrainer internally routes to model A or B
```

### Decision

We choose **Option A**.

Gateway should directly call the specific model’s LLM Server Client handle. MultiModelTrainer should not become an online routing proxy.

Reasons:

1. **Cleaner responsibility boundary**
   Gateway already knows the request context, `AgentID`, session metadata, and routing table. So `AgentID -> ModelID` should happen in Gateway.

2. **Trainer should not be on the online serving path**
   The trainer is responsible for training lifecycle. If all requests go through the trainer, the trainer becomes a serving bottleneck and a new failure point.

3. **Easier to reuse existing Gateway logic**
   We can follow the existing black-box agent / gateway pattern: pass model client handles into Gateway, and let Gateway trigger rollout.

4. **Better isolation**
   If model A is syncing weights, Gateway only needs to pause or retry requests to model A. Model B can still serve normally.

5. **Simpler MultiModelTrainer**
   The trainer only needs to launch runtimes, expose handles, consume trajectories, and manage training queues. It does not need to understand agent workflow routing.

So the final design is:

```text
Gateway does request routing.
Trainer does runtime management and training data ingestion.
```

---

## 4. High-level Architecture

```text
MultiModelDisaggregateTrainer
│
├── ModelAsyncRuntime[model_a]
│   ├── FullyAsyncTrainer[model_a]
│   ├── TrainingQueue[model_a]
│   ├── RolloutRuntime / LLMServerClient[model_a]
│   ├── ParameterSynchronizer[model_a]
│   └── Version / Status Tracker[model_a]
│
├── ModelAsyncRuntime[model_b]
│   ├── FullyAsyncTrainer[model_b]
│   ├── TrainingQueue[model_b]
│   ├── RolloutRuntime / LLMServerClient[model_b]
│   ├── ParameterSynchronizer[model_b]
│   └── Version / Status Tracker[model_b]
│
├── GatewayBindingRegistry
│   ├── model_id -> LLMServerClient
│   ├── model_id -> status client
│   ├── model_id -> version tracker
│   └── agent_id -> model_id
│
└── TrajectoryIngestor
    ├── receives finalized trajectories
    ├── groups by model_id
    └── pushes samples into TrainingQueue[model_id]
```

Request path:

```text
Agent / Workflow
  -> Gateway
  -> AgentID -> ModelID
  -> LLMServerClient[ModelID]
  -> RolloutRuntime[ModelID]
```

Training data path:

```text
Gateway finalized trajectory
  -> TrajectoryIngestor
  -> TrainingQueue[ModelID]
  -> FullyAsyncTrainer[ModelID]
```

---

## 5. Key Components

### 5.1 MultiModelDisaggregateTrainer

The top-level trainer is mainly a launcher and coordinator.

It should:

1. Parse multi-model config.
2. Create one `ModelAsyncRuntime` per model.
3. Start each model’s trainer loop.
4. Start each model’s parameter sync loop.
5. Collect each model’s LLM Server Client handle.
6. Register these handles to Gateway.
7. Start trajectory ingestion.
8. Save per-model checkpoints and global metadata.
9. Report per-model metrics.

It should not:

1. Decide which agent calls which model.
2. Run the agent workflow.
3. Route online generation requests.
4. Block all models when one model is syncing.

Pseudo-code:

```python
class MultiModelDisaggregateTrainer:
    def __init__(self, config):
        self.config = config
        self.runtimes = {}
        self.gateway_registry = None
        self.trajectory_ingestor = None

    async def init(self):
        for model_cfg in self.config.models:
            runtime = ModelAsyncRuntime(model_cfg)
            await runtime.init()
            self.runtimes[model_cfg.model_id] = runtime

        self.gateway_registry = GatewayBindingRegistry(
            agent_to_model=self.config.agent_model_mapping,
            model_clients={
                model_id: runtime.llm_server_client
                for model_id, runtime in self.runtimes.items()
            },
            status_clients={
                model_id: runtime.status_client
                for model_id, runtime in self.runtimes.items()
            },
            version_trackers={
                model_id: runtime.version_tracker
                for model_id, runtime in self.runtimes.items()
            },
        )

        self.trajectory_ingestor = TrajectoryIngestor(
            training_queues={
                model_id: runtime.training_queue
                for model_id, runtime in self.runtimes.items()
            }
        )

    async def fit(self):
        tasks = []

        for runtime in self.runtimes.values():
            tasks.append(asyncio.create_task(runtime.run_trainer_loop()))
            tasks.append(asyncio.create_task(runtime.run_param_sync_loop()))

        tasks.append(asyncio.create_task(self.run_ingestion_loop()))

        await asyncio.gather(*tasks)
```

---

### 5.2 ModelAsyncRuntime

`ModelAsyncRuntime` is a thin wrapper around the existing VERL fully async disaggregate trainer.

Each runtime owns one model’s training and rollout pipeline.

```python
class ModelAsyncRuntime:
    def __init__(self, model_cfg):
        self.model_id = model_cfg.model_id
        self.model_cfg = model_cfg

        self.training_queue = None
        self.trainer = None
        self.rollout_runtime = None
        self.llm_server_client = None
        self.parameter_synchronizer = None
        self.version_tracker = None
        self.status_client = None

    async def init(self):
        self.training_queue = self.create_training_queue()
        self.trainer = self.create_fully_async_trainer(
            queue=self.training_queue,
            model_cfg=self.model_cfg,
        )
        self.rollout_runtime = self.create_rollout_runtime(
            model_cfg=self.model_cfg,
        )
        self.llm_server_client = await self.rollout_runtime.get_client()
        self.parameter_synchronizer = self.create_parameter_synchronizer(
            trainer=self.trainer,
            rollout_runtime=self.rollout_runtime,
        )
        self.version_tracker = VersionTracker(self.model_id)
        self.status_client = RuntimeStatusClient(self.rollout_runtime)

    async def run_trainer_loop(self):
        await self.trainer.run()

    async def run_param_sync_loop(self):
        await self.parameter_synchronizer.run()
```

Implementation note:

Most of this should reuse the existing single-model disaggregate trainer. The main work is to make the single-model trainer initialization reusable instead of hard-coded inside one main training script.

---

### 5.3 GatewayBindingRegistry

This object is the bridge between Trainer and Gateway.

The trainer creates model runtimes and puts their serving handles into the registry. Gateway reads the registry and routes requests.

```python
class GatewayBindingRegistry:
    def __init__(self, agent_to_model, model_clients, status_clients, version_trackers):
        self.agent_to_model = agent_to_model
        self.model_clients = model_clients
        self.status_clients = status_clients
        self.version_trackers = version_trackers

    def get_model_id(self, agent_id: str) -> str:
        return self.agent_to_model[agent_id]

    def get_client(self, model_id: str):
        return self.model_clients[model_id]

    def get_status_client(self, model_id: str):
        return self.status_clients[model_id]

    def get_version_tracker(self, model_id: str):
        return self.version_trackers[model_id]
```

Gateway usage:

```python
model_id = registry.get_model_id(agent_id)
client = registry.get_client(model_id)
response = await client.generate(request)
```

---

### 5.4 Gateway-side Routing

We use the first Gateway design:

```text
One model can have one GatewayManager.
A front router maps AgentID to ModelID.
Then it forwards the request to that model’s GatewayManager / client.
```

Pseudo-code:

```python
class MultiModelGatewayRouter:
    def __init__(self, registry):
        self.registry = registry

    async def chat_completion(self, request):
        agent_id = request.metadata["agent_id"]
        model_id = self.registry.get_model_id(agent_id)

        status_client = self.registry.get_status_client(model_id)
        await status_client.wait_until_ready()

        client = self.registry.get_client(model_id)
        response = await client.generate(request)

        response.metadata["model_id"] = model_id
        response.metadata["agent_id"] = agent_id
        response.metadata["policy_version"] = (
            await self.registry.get_version_tracker(model_id).get_version()
        )

        return response
```

The important point is:

```text
Gateway routes online requests.
Trainer does not route online requests.
```

---

### 5.5 TrajectoryIngestor

`TrajectoryIngestor` is not a request router.

It only handles finalized training data.

```text
trajectory.model_id == model_a -> TrainingQueue[model_a]
trajectory.model_id == model_b -> TrainingQueue[model_b]
```

Pseudo-code:

```python
class TrajectoryIngestor:
    def __init__(self, training_queues):
        self.training_queues = training_queues

    async def ingest(self, trajectories):
        grouped = defaultdict(list)

        for traj in trajectories:
            model_id = traj.metadata["model_id"]
            grouped[model_id].append(traj)

        for model_id, trajs in grouped.items():
            batch = self.convert_to_batch(model_id, trajs)
            await self.training_queues[model_id].put(batch)

    def convert_to_batch(self, model_id, trajectories):
        return DataProto.from_dict(
            {
                "prompt_ids": [t.prompt_ids for t in trajectories],
                "response_ids": [t.response_ids for t in trajectories],
                "response_mask": [t.response_mask for t in trajectories],
                "old_log_probs": [t.old_log_probs for t in trajectories],
                "rewards": [t.reward for t in trajectories],
            },
            meta_info={
                "model_id": model_id,
                "agent_ids": [t.metadata["agent_id"] for t in trajectories],
                "policy_versions": [t.metadata["policy_version"] for t in trajectories],
                "workflow_ids": [t.metadata.get("workflow_id") for t in trajectories],
            },
        )
```

---

## 6. Weight Sync Blocking

### 6.1 Problem

In disaggregate training, the trainer needs to sync updated weights to the rollout side.

During weight sync, the corresponding model’s rollout server may be temporarily unavailable.

In multi-model training, this can happen independently:

```text
model_a is syncing -> model_a rollout may pause
model_b is not syncing -> model_b rollout should still work
```

If a workflow currently needs `model_a`, that workflow may wait. But this should not block `model_b`.

---

### 6.2 Does MultiModelTrainer need special handling?

Mostly no.

The blocking happens when Gateway calls a model whose rollout runtime is syncing. This is a serving-side wait, not a trainer-level global stall.

So MultiModelTrainer does not need a global barrier like:

```text
if one model is syncing, pause all models
```

That would be bad, because it turns local pause into global pause.

The trainer only needs to provide:

1. Per-model runtime status.
2. Per-model policy version.
3. Per-model independent sync.
4. Per-model independent training queue.
5. Per-model independent trainer loop.

Gateway should handle:

1. Waiting for a model to become ready.
2. Retrying a request.
3. Pausing a session.
4. Continuing other sessions that use available models.

---

## 7. Required Code Changes

### 7.1 Trainer side

Add:

```text
multi_model_disaggregate_trainer.py
model_async_runtime.py
gateway_binding_registry.py
trajectory_ingestor.py
checkpoint_coordinator.py
multi_model_metrics.py
```

Main changes:

1. Refactor existing single-model disaggregate trainer initialization into reusable functions.
2. Support multiple `ModelAsyncRuntime`s in one training job.
3. Add per-model queue namespace.
4. Add per-model parameter sync namespace.
5. Add `GatewayBindingRegistry`.
6. Add `TrajectoryIngestor`.
7. Save per-model checkpoints and global metadata.

---

### 7.2 Gateway side

Adopt the first design.

Add or modify:

```text
multi_model_gateway_router.py
agent_model_route_table.py
```

Gateway needs to support:

1. `AgentID -> ModelID`.
2. `ModelID -> LLMServerClient`.
3. Per-model ready status check.
4. Per-request metadata:

   * `agent_id`
   * `model_id`
   * `policy_version`
   * `workflow_id`
5. Finalized trajectory export.

---

### 7.3 Trajectory metadata

Each trajectory should contain:

```python
{
    "workflow_id": str,
    "agent_id": str,
    "model_id": str,
    "policy_version": int,
    "request_id": str,
}
```

This is important for async training, debugging, reward assignment, and checkpoint resume.

---

## 8. Example Config

```yaml
multi_model:
  enable: true

  models:
    - model_id: planner_model
      model_path: /path/to/planner
      trainer:
        n_gpus: 4
      rollout:
        n_gpus: 4
      async_training:
        trigger_parameter_sync_step: 4
        staleness_threshold: 0.5
        partial_rollout: true

    - model_id: solver_model
      model_path: /path/to/solver
      trainer:
        n_gpus: 4
      rollout:
        n_gpus: 4
      async_training:
        trigger_parameter_sync_step: 8
        staleness_threshold: 0.5
        partial_rollout: true

gateway:
  mode: multi_model_router

agent_model_mapping:
  planner_agent: planner_model
  search_agent: solver_model
  answer_agent: solver_model
```

---
---

## 10. Future Optimizations

1. **Stagger weight sync**
   Avoid syncing all models at the same time.

2. **Per-model adaptive concurrency**
   Increase rollout concurrency for models whose training queue is empty.

3. **Per-model batch size**
   Low-frequency models may need smaller training batches.

4. **Per-model staleness threshold**
   Different models can use different async tolerance.

5. **Better reward assignment**
   Start with workflow-level reward broadcast. Later support per-agent or per-turn reward.

6. **Low-frequency model replay buffer**
   Useful when one model is called much less often than others.

---

## 11. Final Summary

The recommended design is:

```text
MultiModelTrainer:
  starts multiple fully async disaggregate runtimes
  manages training queues, parameter sync, checkpoints, metrics
  exposes model-specific LLM Server Client handles

Gateway:
  maps AgentID to ModelID
  directly calls the corresponding LLM Server Client
  handles request-time waiting / retry if a model is syncing

TrajectoryIngestor:
  consumes finalized trajectories
  pushes samples to the correct model’s training queue
```

We should not put request routing inside MultiModelTrainer.

The trainer should be a training runtime manager, not an online serving router.

---
## 1. 提案

本 RFC 提议实现一个 `MultiModelDisaggregateTrainer`，用于在 Uni-Agent / VERL 上支持 multi-model co-training。

核心思路很简单：

```text
一个模型 = 一个 fully async disaggregate runtime。
MultiModelDisaggregateTrainer = runtime 启动器 + 生命周期管理器 + 训练数据消费者。
Gateway = 请求路由器。
```

每个模型尽量复用 VERL 现有的 fully async disaggregate trainer。我们不从零写一个新 trainer，而是把每个 single-model async trainer 包成一个 `ModelAsyncRuntime`，然后由外层 `MultiModelDisaggregateTrainer` 同时启动和管理多个 runtime。

对于模型调用，我们采用 **Gateway 侧路由**：

```text
Agent request
  -> Gateway
  -> AgentID -> ModelID
  -> 对应模型的 LLM Server Client
```

Multi-model Trainer 不应该返回一个统一的 LLM Server Client，然后在 trainer 内部再做 route。Trainer 应该负责启动各个模型 runtime，拿到每个模型的 LLM Server Client 句柄，并把这些句柄注册给 Gateway。Gateway 再直接调用对应模型的 client。

---

## 2. 动机

### 2.1 提升 rollout 效率

在 colocate 模式下，训练和 rollout 绑得很紧。多模型训练时，不同模型的训练速度可能不同。比如 model A 很快 train 完了，但 model B 还在做 optimizer step。此时即使 model A 已经 ready，整个 agent task 也可能没法继续 rollout，因为另一个模型还在训练。

Disaggregate 模式下，训练和 rollout 分离。除了短暂的 weight sync 时间以外，rollout server 大多数时间都可以服务请求。这对 multi-agent workload 会更高效。

```text
Colocate:
  所有模型一起 train -> 等最慢的模型 -> rollout

Disaggregate:
  每个模型独立 train
  只要该模型不在 sync，就可以 rollout
```

### 2.2 缓解长尾效应

Colocate 模式本来就有长尾问题：有些 rollout sample 很慢，整个 batch 要等它。

在 multi-agent 场景下，长尾会更严重，因为不同 agent / model 的长尾可能叠加。一个慢模型调用，或者一个慢 agent trajectory，都可能拖住整个 pipeline。

Async + disaggregate 可以缓解这个问题，因为 rollout 和 training 变成 streaming-style，不同模型可以独立推进，不需要每一步都互相等。

### 2.3 更适合 dynamic agent workflow

在 dynamic agent workflow 里，每个模型被调用的次数可能完全不同。

比如：

```text
planner model: 调一次
tool-use model: 调很多次
verifier model: 只在困难样本上调用
```

这些模型凑 batch 的速度不一样。如果全部 colocate，不同模型会频繁互相等待，这不太合理。

Disaggregate 更自然。每个模型都有自己的 rollout service、training queue、trainer 和 sync schedule。高频模型可以更频繁训练，低频模型可以用更小 batch 或更低 update frequency。

---

## 3. 架构选择：Gateway 直接调用模型 Client

这里有两种设计。

### 方案 A：Gateway 直接调用每个模型的 LLM Server Client

```text
MultiModelTrainer
  -> 启动 ModelRuntime[A]
  -> 启动 ModelRuntime[B]
  -> 拿到 LLMServerClient[A], LLMServerClient[B]
  -> 注册给 Gateway

Gateway
  -> 收到 AgentID
  -> 映射到 ModelID
  -> 调用 LLMServerClient[ModelID]
```

### 方案 B：Gateway 只调用 MultiModelTrainer 暴露的统一 Client

```text
Gateway
  -> 调用 UnifiedMultiModelClient
  -> MultiModelTrainer 内部再 route 到 model A 或 model B
```

### 结论

我们选择 **方案 A**。

Gateway 应该直接调用某个具体模型的 LLM Server Client 句柄。MultiModelTrainer 不应该变成在线请求路由代理。

原因：

1. **职责更清晰**
   Gateway 知道请求上下文、`AgentID`、session metadata 和 routing table。所以 `AgentID -> ModelID` 应该由 Gateway 做。

2. **Trainer 不应该在在线服务链路上**
   Trainer 负责 training lifecycle。如果所有请求都先打到 trainer，trainer 会变成 serving bottleneck 和新的 failure point。

3. **更容易复用现有 Gateway 逻辑**
   可以沿用现在 black-box agent / gateway 的模式：把模型 client handle 传给 Gateway，由 Gateway 触发 rollout。

4. **隔离性更好**
   如果 model A 正在 sync weight，Gateway 只需要暂停或 retry model A 的请求。model B 仍然可以正常服务。

5. **MultiModelTrainer 更简单**
   Trainer 只需要启动 runtime、暴露 handle、消费 trajectory、管理训练队列。它不需要理解 agent workflow routing。

最终边界是：

```text
Gateway 做请求路由。
Trainer 做 runtime 管理和训练数据 ingestion。
```

---

## 4. 总体架构

```text
MultiModelDisaggregateTrainer
│
├── ModelAsyncRuntime[model_a]
│   ├── FullyAsyncTrainer[model_a]
│   ├── TrainingQueue[model_a]
│   ├── RolloutRuntime / LLMServerClient[model_a]
│   ├── ParameterSynchronizer[model_a]
│   └── Version / Status Tracker[model_a]
│
├── ModelAsyncRuntime[model_b]
│   ├── FullyAsyncTrainer[model_b]
│   ├── TrainingQueue[model_b]
│   ├── RolloutRuntime / LLMServerClient[model_b]
│   ├── ParameterSynchronizer[model_b]
│   └── Version / Status Tracker[model_b]
│
├── GatewayBindingRegistry
│   ├── model_id -> LLMServerClient
│   ├── model_id -> status client
│   ├── model_id -> version tracker
│   └── agent_id -> model_id
│
└── TrajectoryIngestor
    ├── 接收 finalized trajectories
    ├── 按 model_id 分组
    └── 推到 TrainingQueue[model_id]
```

请求路径：

```text
Agent / Workflow
  -> Gateway
  -> AgentID -> ModelID
  -> LLMServerClient[ModelID]
  -> RolloutRuntime[ModelID]
```

训练数据路径：

```text
Gateway finalized trajectory
  -> TrajectoryIngestor
  -> TrainingQueue[ModelID]
  -> FullyAsyncTrainer[ModelID]
```

---

## 5. 核心组件

### 5.1 MultiModelDisaggregateTrainer

外层 trainer 主要是 launcher 和 coordinator。

它需要做：

1. 解析 multi-model config。
2. 给每个 model 创建一个 `ModelAsyncRuntime`。
3. 启动每个 model 的 trainer loop。
4. 启动每个 model 的 parameter sync loop。
5. 收集每个 model 的 LLM Server Client 句柄。
6. 把这些句柄注册给 Gateway。
7. 启动 trajectory ingestion。
8. 保存 per-model checkpoint 和 global metadata。
9. 记录 per-model metrics。

它不应该做：

1. 决定哪个 agent 调哪个模型。
2. 执行 agent workflow。
3. 路由在线 generation request。
4. 一个模型 sync 时 block 所有模型。

伪代码：

```python
class MultiModelDisaggregateTrainer:
    def __init__(self, config):
        self.config = config
        self.runtimes = {}
        self.gateway_registry = None
        self.trajectory_ingestor = None

    async def init(self):
        for model_cfg in self.config.models:
            runtime = ModelAsyncRuntime(model_cfg)
            await runtime.init()
            self.runtimes[model_cfg.model_id] = runtime

        self.gateway_registry = GatewayBindingRegistry(
            agent_to_model=self.config.agent_model_mapping,
            model_clients={
                model_id: runtime.llm_server_client
                for model_id, runtime in self.runtimes.items()
            },
            status_clients={
                model_id: runtime.status_client
                for model_id, runtime in self.runtimes.items()
            },
            version_trackers={
                model_id: runtime.version_tracker
                for model_id, runtime in self.runtimes.items()
            },
        )

        self.trajectory_ingestor = TrajectoryIngestor(
            training_queues={
                model_id: runtime.training_queue
                for model_id, runtime in self.runtimes.items()
            }
        )

    async def fit(self):
        tasks = []

        for runtime in self.runtimes.values():
            tasks.append(asyncio.create_task(runtime.run_trainer_loop()))
            tasks.append(asyncio.create_task(runtime.run_param_sync_loop()))

        tasks.append(asyncio.create_task(self.run_ingestion_loop()))

        await asyncio.gather(*tasks)
```

---

### 5.2 ModelAsyncRuntime

`ModelAsyncRuntime` 是 VERL 现有 fully async disaggregate trainer 的薄封装。

每个 runtime 拥有一个模型的 training 和 rollout pipeline。

```python
class ModelAsyncRuntime:
    def __init__(self, model_cfg):
        self.model_id = model_cfg.model_id
        self.model_cfg = model_cfg

        self.training_queue = None
        self.trainer = None
        self.rollout_runtime = None
        self.llm_server_client = None
        self.parameter_synchronizer = None
        self.version_tracker = None
        self.status_client = None

    async def init(self):
        self.training_queue = self.create_training_queue()
        self.trainer = self.create_fully_async_trainer(
            queue=self.training_queue,
            model_cfg=self.model_cfg,
        )
        self.rollout_runtime = self.create_rollout_runtime(
            model_cfg=self.model_cfg,
        )
        self.llm_server_client = await self.rollout_runtime.get_client()
        self.parameter_synchronizer = self.create_parameter_synchronizer(
            trainer=self.trainer,
            rollout_runtime=self.rollout_runtime,
        )
        self.version_tracker = VersionTracker(self.model_id)
        self.status_client = RuntimeStatusClient(self.rollout_runtime)

    async def run_trainer_loop(self):
        await self.trainer.run()

    async def run_param_sync_loop(self):
        await self.parameter_synchronizer.run()
```

实现上，应该尽量复用现有 single-model disaggregate trainer。主要改动是把 single-model trainer 的初始化逻辑抽出来，不要写死在一个 main training script 里。

---

### 5.3 GatewayBindingRegistry

这个对象负责连接 Trainer 和 Gateway。

Trainer 创建 model runtimes 后，把每个模型的 serving handle 放到 registry 里。Gateway 从 registry 里读取 handle 并路由请求。

```python
class GatewayBindingRegistry:
    def __init__(self, agent_to_model, model_clients, status_clients, version_trackers):
        self.agent_to_model = agent_to_model
        self.model_clients = model_clients
        self.status_clients = status_clients
        self.version_trackers = version_trackers

    def get_model_id(self, agent_id: str) -> str:
        return self.agent_to_model[agent_id]

    def get_client(self, model_id: str):
        return self.model_clients[model_id]

    def get_status_client(self, model_id: str):
        return self.status_clients[model_id]

    def get_version_tracker(self, model_id: str):
        return self.version_trackers[model_id]
```

Gateway 使用方式：

```python
model_id = registry.get_model_id(agent_id)
client = registry.get_client(model_id)
response = await client.generate(request)
```

---

### 5.4 Gateway 侧路由

采用第一种 Gateway 方案：

```text
每个模型可以有自己的 GatewayManager。
前面加一个 front router，根据 AgentID 找 ModelID。
然后把请求转发给对应模型的 GatewayManager / client。
```

伪代码：

```python
class MultiModelGatewayRouter:
    def __init__(self, registry):
        self.registry = registry

    async def chat_completion(self, request):
        agent_id = request.metadata["agent_id"]
        model_id = self.registry.get_model_id(agent_id)

        status_client = self.registry.get_status_client(model_id)
        await status_client.wait_until_ready()

        client = self.registry.get_client(model_id)
        response = await client.generate(request)

        response.metadata["model_id"] = model_id
        response.metadata["agent_id"] = agent_id
        response.metadata["policy_version"] = (
            await self.registry.get_version_tracker(model_id).get_version()
        )

        return response
```

重点是：

```text
Gateway 路由在线请求。
Trainer 不路由在线请求。
```

---

### 5.5 TrajectoryIngestor

`TrajectoryIngestor` 不是 request router。

它只处理已经完成的训练数据。

```text
trajectory.model_id == model_a -> TrainingQueue[model_a]
trajectory.model_id == model_b -> TrainingQueue[model_b]
```

伪代码：

```python
class TrajectoryIngestor:
    def __init__(self, training_queues):
        self.training_queues = training_queues

    async def ingest(self, trajectories):
        grouped = defaultdict(list)

        for traj in trajectories:
            model_id = traj.metadata["model_id"]
            grouped[model_id].append(traj)

        for model_id, trajs in grouped.items():
            batch = self.convert_to_batch(model_id, trajs)
            await self.training_queues[model_id].put(batch)

    def convert_to_batch(self, model_id, trajectories):
        return DataProto.from_dict(
            {
                "prompt_ids": [t.prompt_ids for t in trajectories],
                "response_ids": [t.response_ids for t in trajectories],
                "response_mask": [t.response_mask for t in trajectories],
                "old_log_probs": [t.old_log_probs for t in trajectories],
                "rewards": [t.reward for t in trajectories],
            },
            meta_info={
                "model_id": model_id,
                "agent_ids": [t.metadata["agent_id"] for t in trajectories],
                "policy_versions": [t.metadata["policy_version"] for t in trajectories],
                "workflow_ids": [t.metadata.get("workflow_id") for t in trajectories],
            },
        )
```

---

## 6. Weight Sync 阻断问题

### 6.1 问题

Disaggregate training 里，trainer 更新权重后，需要把新权重同步到 rollout 侧。

同步期间，对应模型的 rollout server 可能会短暂不可用。

在 multi-model training 里，这个事情是 per-model 独立发生的：

```text
model_a 正在 sync -> model_a rollout 可能暂停
model_b 没有 sync -> model_b rollout 应该继续可用
```

如果某个 workflow 当前需要 `model_a`，那这个 workflow 可能要等。但这不应该 block `model_b`。

---

### 6.2 MultiModelTrainer 需要特殊处理吗？

基本不需要。

阻断发生在 Gateway 调用某个正在 sync 的模型时。这是 serving 侧的等待，不是 trainer 侧的全局停顿。

所以 MultiModelTrainer 不需要做这种 global barrier：

```text
只要有一个模型在 sync，就暂停所有模型
```

这会把局部 pause 扩大成全局 pause，反而违背 fully async disaggregate 的目标。

Trainer 只需要提供：

1. per-model runtime status；
2. per-model policy version；
3. per-model independent sync；
4. per-model independent training queue；
5. per-model independent trainer loop。

Gateway 负责：

1. 等待模型 ready；
2. retry 请求；
3. 暂停 session；
4. 继续执行其他使用可用模型的 session。

---

## 7. 需要修改的代码

### 7.1 Trainer 侧

新增：

```text
multi_model_disaggregate_trainer.py
model_async_runtime.py
gateway_binding_registry.py
trajectory_ingestor.py
checkpoint_coordinator.py
multi_model_metrics.py
```

主要改动：

1. 把现有 single-model disaggregate trainer 初始化逻辑抽成可复用函数。
2. 支持一个 training job 里启动多个 `ModelAsyncRuntime`。
3. 增加 per-model queue namespace。
4. 增加 per-model parameter sync namespace。
5. 增加 `GatewayBindingRegistry`。
6. 增加 `TrajectoryIngestor`。
7. 保存 per-model checkpoint 和 global metadata。

---

### 7.2 Gateway 侧

采用第一种设计。

新增或修改：

```text
multi_model_gateway_router.py
agent_model_route_table.py
```

Gateway 需要支持：

1. `AgentID -> ModelID`。
2. `ModelID -> LLMServerClient`。
3. per-model ready status check。
4. per-request metadata：

   * `agent_id`
   * `model_id`
   * `policy_version`
   * `workflow_id`
5. finalized trajectory export。

---

### 7.3 Trajectory metadata

每条 trajectory 需要包含：

```python
{
    "workflow_id": str,
    "agent_id": str,
    "model_id": str,
    "policy_version": int,
    "request_id": str,
}
```

这些字段对 async training、debug、reward assignment 和 checkpoint resume 都很重要。

---

## 8. 配置示例

```yaml
multi_model:
  enable: true

  models:
    - model_id: planner_model
      model_path: /path/to/planner
      trainer:
        n_gpus: 4
      rollout:
        n_gpus: 4
      async_training:
        trigger_parameter_sync_step: 4
        staleness_threshold: 0.5
        partial_rollout: true

    - model_id: solver_model
      model_path: /path/to/solver
      trainer:
        n_gpus: 4
      rollout:
        n_gpus: 4
      async_training:
        trigger_parameter_sync_step: 8
        staleness_threshold: 0.5
        partial_rollout: true

gateway:
  mode: multi_model_router

agent_model_mapping:
  planner_agent: planner_model
  search_agent: solver_model
  answer_agent: solver_model
```

---

## 10. 后续优化

1. **错开 weight sync**
   避免所有模型同时 sync。

2. **Per-model adaptive concurrency**
   如果某个模型 training queue 太空，可以增加该模型相关 rollout 并发。

3. **Per-model batch size**
   低频模型可以用更小 training batch。

4. **Per-model staleness threshold**
   不同模型可以有不同 async tolerance。

5. **更细粒度 reward assignment**
   第一版可以 workflow-level reward broadcast。后续支持 per-agent 或 per-turn reward。

6. **低频模型 replay buffer**
   如果某个模型调用次数很少，可以引入 replay buffer。

---

## 11. 总结

推荐设计是：

```text
MultiModelTrainer:
  启动多个 fully async disaggregate runtimes
  管理 training queue、parameter sync、checkpoint、metrics
  暴露每个模型自己的 LLM Server Client handle

Gateway:
  根据 AgentID 找 ModelID
  直接调用对应模型的 LLM Server Client
  如果模型正在 sync，在 request-time 做等待 / retry

TrajectoryIngestor:
  消费 finalized trajectories
  把 sample 推到对应模型的 training queue
```

不要把 request routing 放进 MultiModelTrainer。

Trainer 应该是 training runtime manager，不应该是 online serving router。


### Motivation and use case

stated above

### Related area

core agent loop

### Proposed design or API

_No response_

### Alternatives considered

_No response_

### Additional context

_No response_

### Contribution

- [ ] I am willing to help implement this feature.
- [ ] I am willing to help test or validate this feature.

Uh oh!

[RFC]feat: Multi-policy Disaggregate Trainer #77

Description

RFC: Multi-policy (models) Fully Async Disaggregate Trainer

1. Proposal

2. Motivation

2.1 Improve rollout efficiency

2.2 Reduce long-tail effect

2.3 Better fit for dynamic agent workflows

3. Architecture Decision: Gateway Routes Directly to Model Clients

Option A: Gateway directly calls each model’s LLM Server Client

Option B: Gateway calls one unified client from MultiModelTrainer

Decision

4. High-level Architecture

5. Key Components

5.1 MultiModelDisaggregateTrainer

5.2 ModelAsyncRuntime

5.3 GatewayBindingRegistry

5.4 Gateway-side Routing

5.5 TrajectoryIngestor

6. Weight Sync Blocking

6.1 Problem

6.2 Does MultiModelTrainer need special handling?

7. Required Code Changes

7.1 Trainer side

7.2 Gateway side

7.3 Trajectory metadata

8. Example Config

10. Future Optimizations

11. Final Summary

1. 提案

2. 动机

2.1 提升 rollout 效率

2.2 缓解长尾效应

2.3 更适合 dynamic agent workflow

3. 架构选择：Gateway 直接调用模型 Client

方案 A：Gateway 直接调用每个模型的 LLM Server Client

方案 B：Gateway 只调用 MultiModelTrainer 暴露的统一 Client

结论

4. 总体架构

5. 核心组件

5.1 MultiModelDisaggregateTrainer

5.2 ModelAsyncRuntime

5.3 GatewayBindingRegistry

5.4 Gateway 侧路由

5.5 TrajectoryIngestor

6. Weight Sync 阻断问题

6.1 问题

6.2 MultiModelTrainer 需要特殊处理吗？

7. 需要修改的代码

7.1 Trainer 侧

7.2 Gateway 侧

7.3 Trajectory metadata

8. 配置示例

10. 后续优化

11. 总结

Motivation and use case

Related area

Proposed design or API

Alternatives considered

Additional context

Contribution

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions