Remove renderer from Rollouter group API - #4076
Conversation
meta-cla
Bot
added
the
CLA Signed
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #4076 * #4074 * #4073 * __->__ #4072 `InterGeneratorRouter` has mutable states such as generator load and lifecycle state. Its current implementation is not safe to share across threads or processes. This blocks it from being used concurrently from multiple processes. To solve this problem, this PR add a `InterGeneratorRouterActor`, which is a Monarch Actor wrapper. This actor is spawned on its dedicate process as a singleton Monarch actor mesh. In this way, the router can be accessed concurrently through this actor. Currently the actor is only used by controller. This PR is mainly to set up the stage so we later can: 1. Put Rollouters on a process pool, where each process needs its own router access. 2. Put router behind a HTTP endpoint, where the HTTP server is in its own process, or maybe even a different host.
|
|
||
| await self._rollouter.setup_async() | ||
| await self._rollouter.setup_async( | ||
| renderer_config=config.renderer, |
There was a problem hiding this comment.
Sounds safe and no chance of breaking token-in-token-out -- the only other use case is to get the eos_token_id for batcher to do padding https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/rl/controller.py#L801
but we don't have to do that -- for attention the masks are built using positions; for moe all padding tokens will be routed to the same experts. So it sounds we can use anything? cc @acisseJZhong
The main use case in TokenEnv
https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/rl/environment/token.py#L109
But since it's not supposed to be user facing, maybe we can put renderer config into RolloutWorker.
There was a problem hiding this comment.
not sure if we should ignore this eos_token_id and pad with random number(you mentioned there might be issue with deterministic RNG state?), or add mask to MoE forward. I need to think more.
regardless of which approach, after this PR we can still access renderer eos_token_id in controller right?
There was a problem hiding this comment.
we can, but I'm like "if we can remove access of renderer in controller, we can make it cleaner and pleasing; and it doesn't sound absolutely necessary for controller to access it", lol
| group_id: int, | ||
| group_size: int, | ||
| sampling: SamplingConfig, | ||
| renderer: Renderer, |
There was a problem hiding this comment.
This IPC overhead is expensive for large payload such as renderer
curious how did you identify this? Experience or by looking at logs / traces
There was a problem hiding this comment.
I forgot it was me or claude spotting it first. But it was not from logs.
In production, issues like this should be detected by the "endpoint payload" metric, which will be added to monarch.
3 participants
Stack from ghstack (oldest at bottom):
After #4074,
run_group_rollouts's underlying logic is executed in process different from the controller process. This mean callingrun_group_rolloutsfrom controller will result in serializing all its parameters, includerenderer, and then sent to a different process. This IPC overhead is expensive for large payload such asrenderer, so we need to remove it from the parameter list, but make it a field of the worker. In that way,rendererwill only be sent over IPC once duringset_async.However, this change means it forfeits the flexibility of injecting
Rendererintomake_env_group:torchtitan/torchtitan/experiments/rl/rollout/rollouter.py
Lines 111 to 113 in ecae62f
The need for injection is based on the assumption that we might have multiple
Renderers. In that case, we could chose to index theseRenderers, and only pass their IDs from controller. However, I do not think we need to do that now.