[Feature Request] Auto-batching inference server

Motivation

I am always disappointed when my GPUs are underutilized. I am currently working on having several agents operate in parallel on an MCTS tree. As the actors explore the tree independently, they need to query the model to select the next action. While each actor could technically keep its own copy of the model, this does not scale to a large number of parallel actors and can underutilize GPU resources.

Ideally, we would batch all model calls from the actors and serve them using one or more GPUs. In the LLM setting, this can be achieved by running a vLLM server: each actor sends requests to a shared server, which handles automatic batching, etc. While there are alternatives (e.g., Triton Inference Server or TorchServe, though TorchServe appears to be in limited maintenance), these are not ideal in a training loop.

The main issue is that there is no graceful way to update the parameters of the served model (and no tensordicts). With Triton, for example, you typically need to save a checkpoint, place it into the model repository, and then trigger a model load/reload. While workable, this is not elegant, and it can be slow enough that frequent updates become prohibitive.

Solution

We can build our own more general inference server, potentially giving up some performance in exchange for a more training-friendly workflow. The properties we want are automatic batching, returning TensorDicts, a custom collate_fn, and the ability to integrate model weights with torch.distributed.

Queue
The queue is the primary interface for interacting with the served model. The inference worker consumes up to max_batch elements from the queue (when enough elements are available), or consumes all currently available elements once a time_limit is reached.

TensorDicts
The inference server should return the model’s TensorDict outputs to the appropriate actor. I’m not sure what the best way to implement this routing is.

collate_fn
Different users have different data formats, and variable-length data may require packing or padding. It would be nice to keep this flexible.

Weight sharing
We place the served model and the training model in a shared distributed group. During training, the user can broadcast updated weights efficiently to the inference model to minimize time spent waiting on synchronization.

Alternatives

I considered Triton Inference Server and TorchServe, but neither satisfies the requirements above.

Additional context

N/A

Checklist

• I have checked that there is no similar issue in the repo (required)