Investigate on adding new MLX roll out engines for Verl

[nguyenhoangthuan99]

Literature review https://github.com/volcengine/verl/tree/main/verl/workers/rollout

Currently support vLLM and SGLang, if we need to add, what APIs we have to implement?

[nguyenhoangthuan99]

How to implement new Rollout method

• Registry that maps (name, mode) -> class FQDN: see verl/workers/rollout/base.py:80 — the registry currently contains vllm and sglang entries.
• SGLang rollout example (full-featured sync + server adapters): class starts at verl/workers/rollout/sglang_rollout/sglang_rollout.py:249 and the ServerAdapter at verl/workers/rollout/sglang_rollout/sglang_rollout.py:1528.
• vLLM rollout example (sync spmd + async wrapper): class vLLMRollout at verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:92 and the async adapter at verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:472.
• Utilities for networking and server startup in verl/workers/rollout/utils.py:26.

How rollout selection works (quick)

• The function get_rollout_class in verl/workers/rollout/base.py:88 looks up the tuple (rollout_name, mode) in the module-level _ROLLOUT_REGISTRY and import/getattr the class. So adding a new rollout consists of:
  1. implementing a BaseRollout subclass (and optional server adapter),
  2. registering its FQDN in _ROLLOUT_REGISTRY under the chosen keys.

What a rollout class must provide

• All rollout implementations inherit from BaseRollout (verl/workers/rollout/base.py:28).
• Required abstract methods (signatures you must implement):
  • async def resume(self, tags: list[str])
  • async def update_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None], **kwargs)
  • async def release(self)
  • def generate_sequences(self, prompts: DataProto) -> DataProto (sync path — for server/async variants the runtime may call different adapters)
• You’ll also usually accept the same constructor signature:
  def init(self, config: RolloutConfig, model_config: HFModelConfig, device_mesh: DeviceMesh)
  (see BaseRollout init at verl/workers/rollout/base.py:31).

Concrete step-by-step recipe to add a new rollout "myrollout"

1. Create new package and file

• Add a directory and file, e.g.:
  • file: verl/workers/rollout/myrollout/myrollout.py:1
• Implement two classes (if you want both sync and async modes):
  • MyRollout (sync / SPMD-like / in-process)
  • MyRolloutServerAdapter (async / server-client style)

2. Register your rollout class

• Edit registry map in verl/workers/rollout/base.py:80.

• Add entries for your rollout and both modes you support. Example patch lines to add:

  _ROLLOUT_REGISTRY = {
  ...
  ("myrollout", "sync"): "verl.workers.rollout.myrollout.myrollout.MyRollout",
  ("myrollout", "async"): "verl.workers.rollout.myrollout.myrollout.MyRolloutServerAdapter",
  }

• After adding to _ROLLOUT_REGISTRY, code that calls get_rollout_class("myrollout", "sync") will import your class automatically.

3. Decide API details and config flags

• Rollout classes typically read items from RolloutConfig and HFModelConfig passed to the constructor.
  • Example keys seen in SGLang/vLLM: load_format, dtype, free_cache_engine, tensor_model_parallel_size, response_length, prompt_length, calculate_log_probs, update_weights_bucket_megabytes.
• If you need new config options, add them into the project's rollout config or pass via existing RolloutConfig (see files under verl/workers/config/ for patterns).
• The rest of the code (trainer/actor) uses get_rollout_class(name, mode) and instantiates your rollout with these configs.

4. Implement weight updates efficiently

• update_weights receives a generator of (name, tensor). Large models require bucketization — see SGLang's use of get_named_tensor_buckets and sgl_update_weights in verl/workers/rollout/sglang_rollout/sglang_rollout.py:1509.
• If your engine provides a fast bulk loader (DTensor/DT loader), use that; else implement bucket splitting to avoid huge temporary allocations.

5. Implement resume/release for free-cache behavior

• Many rollout engines support “sleep/wake” (free cache engine) to release GPU memory when not used; implement resume/release accordingly (see SGLang resume/release at verl/workers/rollout/sglang_rollout/sglang_rollout.py:1488 and vLLM’s wake_up/sleep at verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:407).

6. Tokenizer / processor handling

• If you need tokenizer or a Processor (for multimodal), get it from model_config:
  • In SGLang they call model_config.get_processor() and store pad_token_id: see verl/workers/rollout/sglang_rollout/sglang_rollout.py:258.
  • vLLM rollout uses model_config.tokenizer (verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:107).

Minimal checklist to implement (order)

• Create new package verl/workers/rollout/myrollout/myrollout.py:1 and implement MyRollout class with required methods.
• Optionally implement server adapter class for async mode.
• Register entries in registry at verl/workers/rollout/base.py:80.
• Add config knobs if needed (follow patterns in verl/workers/config/).
• Implement efficient update_weights handling (buckets or engine bulk loader).
• Run a local smoke test by instantiating via get_rollout_class and calling generate_sequences with a small DataProto.

[nguyenhoangthuan99]

Implement MLX engine like vLLM

Short answer: vLLM is used as an inference engine with a small surface the rollout code expects. To implement your own engine you must provide an equivalent API surface (constructor + a few methods / runtime objects) and matching runtime behavior (batched generation, sleep/wake, weight loading, optional LoRA injection, an async worker wrapper if you want server mode). Below I list the vLLM APIs the code calls (with exact call sites) and then a concrete checklist of what you must implement to act like vLLM for this project.

What the code calls from vLLM (direct references)

• Engine construction: the code instantiates the vLLM engine via LLM(...). This is the main engine object the rollout uses.
• Synchronous generation call: rollout calls inference_engine.generate(...) and then iterates the returned outputs to access:
  • each sample’s token ids via .outputs[sample_id].token_ids and
  • each sample’s logprob entries via .outputs[sample_id].logprobs
    (see the loop around [verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:350]).
• Sleep / wake (free-cache behavior): rollout calls inference_engine.wake_up(tags=...) and inference_engine.sleep(level=...) to manage GPU residency.
• Prefix-cache reset: rollout calls [inference_engine.reset_prefix_cache()] to clear prefix cache before sleeping (see [verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:422]).
• LoRA runtime API:
  • list existing LoRAs: [inference_engine.llm_engine.list_loras()] (used at [verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:331])
  • add a LoRA request at runtime: [inference_engine.llm_engine.add_lora(...)] (used at [verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:446]).
• Weight loading for update_weights: when not using PEFT flow, the code obtains the underlying model object and calls a loader (model.load_weights(weights)) after applying a vLLM MOE patch — see [verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:451], and later [verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:603].
• Async server wrapper pattern: the async version exposes a ZeroMQ endpoint that receives RPCs and delegates to methods like init_worker, load_model, and execute_method on a WorkerWrapperBase (see [verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py:537 → :566]). The rollout’s ZeroMQ loop deserializes (method, args, kwargs) and calls that method on the worker.