[Core] Refactor Worker and ModelRunner to consolidate control plane communication

[stephanie-wang]

Currently the ModelRunner class contains both model execution code and multi-GPU control plane communication code, i.e.broadcast_tensor_dict calls. This makes it difficult to swap out the control plane mechanism, e.g., using NCCL vs CPU-based serialization to move the inputs from the LLMEngine to the Workers. It also makes it difficult to improve performance for multi-GPU settings, because we need to reason about all of the possible broadcast calls that may happen during one execution request.

Previously, the Worker and ModelRunner class looked something like this:

class Worker:
  def execute_model(self, execute_model_req=None):
    # Do some broadcast here.
    ...
    return self.model_runner.execute_model(execute_model_req)

class ModelRunner:
  def execute_model(self, execute_model_req=None):
    # Do some more broadcast here.
    ...
    return model_executable(...)

This PR refactors the Worker and ModelRunner classes to consolidate all of the control plane communication into a central location. This adds a new class ModelInput, which acts as the interface between Worker and ModelRunner. The contract is that ExecuteModelRequest contains CPU-only metadata, while ModelInput's tensors should already be on the correct device. Now, the Worker class looks like this:

class Worker:
  def execute_model(self, execute_model_req=None):
    # *All* broadcast / other communication happens in this function.
    model_input: ModelInput = self.prepare_model_input(execute_model_req)
    # Worker.execute_model_local and ModelRunner code contain single-device code only.
    return self.execute_model_local(model_input)

This refactor has a few advantages:

• Reduces the number of broadcast_tensor_dict calls, which can save some latency
• Cleaner separation between control plane communication vs. single-GPU logic. Each ModelRunner needs to explicitly state what inputs it requires by defining a ModelInput (subclass). This requires a bit more developer effort but should make it easier to try out other optimizations in the future:
  • Other control plane implementations (e.g., CPU-based serialization instead of NCCL)
  • Moving the rank 0 worker to a different process from the LLMEngine.
  • Pipelining model execution with metadata movement between CPU and GPU

One complication is that the speculative decoding code goes back and forth between ExecuteModelRequest and ModelInput, whereas other workers only convert from ExecuteModelRequest to ModelInput. Thus, for the speculative decoding path, it's easier for now to keep the per-step broadcast. These extra k broadcasts could also be consolidated in the future, by either supporting ModelInput -> ExecuteModelRequest, or by making it possible to modify a ModelInput.

TODO:

• Update neuron model runner
• Disallow control plane communication ops inside ModelRunner code

[youkaichao]

Thanks, Stephanie! FYI I have a PR with somewhat similar functionality #5354 .

[stephanie-wang]

Thanks, Stephanie! FYI I have a PR with somewhat similar functionality #5354 .

Ah I see, thanks! Do you want to take a quick look at this PR and see what you think? I think this is a superset of #5354.

[youkaichao]

Will take a look! I like this idea:

The contract is that ExecuteModelRequest contains CPU-only metadata, while ModelInput's tensors should already be on the correct device.

[rkooo567]

Looks pretty clean! A little worried about spec decoding case, but I guess there must be a solution.

IIUC, we cannot run tests with execute_model_local with this PR right? Are we planning to add test in the future PRs?

[rkooo567]

added action-required flag (lmk when it is ready again), and cc @cadedaniel. I think after my comments are addressed, tests pass, I am pretty happy with a PR now

[youkaichao]

why do we need this? originally we already have .model_runner attribute.

[youkaichao]

we can add type annotation here:

    is_driver_worker: bool
    model_runner: Typexxx

Then type check will know these attributes. And we just need to let some workers inherit this class, without modifying their self.is_driver_worker to self._is_driver_worker .

[youkaichao]

LGTM, thanks for addressing all my comments!