[Tinker] Refactor trainer and worker (to move algo to trainer and infra to worker)

[tyler-griggs]

PR Overview

This PR has a pretty significant architectural refactor. The primary changes are:

1. Move algorithm logic into the trainer and infra logic into the workers.

   • Remove ppo_train (from FSDP workers for now, @erictang000 to remove from Megatron)
   • Move micro-batching logic from trainer to worker (worker doesn't know mini batches, trainer doesn't know micro batches
   • Worker onloading/offloading moved out of trainer

2. Introduce WorkerDispatch: an intermediate layer between the training loop and workers. Later, this is the entity that will sit beneath the Tinker API server. It is the entity that handles a "pool" of workers. Currently it dispatches requests to the pool of workers (by mesh or by pass_through) and handles onloading / offloading workers.
   i. I'm very open to renaming suggestions.

---

Architecture

Here's a diagram of what the breakdown of responsibilities now looks like. Note that some of this will still need to change. E.g., the trainer will not call dispatch in the future, it will call into the Tinker API server.

┌─────────────────────────────────────────────────────────────┐
│                    TRAINER (Algorithm)                       │
│  - PPO algorithm implementation                              │
│  - Knows only mini batches                                   │
│  - Calls dispatch.forward_backward() + dispatch.optim_step() │
└─────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────┐
│                 WORKER DISPATCH (Coordination)               │
│  - Manages all actor groups (policy, critic, ref)           │
│  - Handles GPU state (offload/backload) automatically       │
│  - Routes calls to appropriate workers                      │
│  - Handles DP sharding via MeshDispatch                     │
└─────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────┐
│                    WORKERS (Execution)                       │
│  - Execute forward/backward passes                          │
│  - Handle micro-batching internally                         │
│  - Scale gradients at optim_step                            │
│  - Model-specific implementations (FSDP, Megatron)          │
└─────────────────────────────────────────────────────────────┘

---

Tests

Deleted Tests

• test_ppo_train.py - tested removed ppo_train method

Updated Tests

• test_training_step.py - uses WorkerDispatch for policy tests
• test_worker_offload.py - updated to work with new interfaces
• test_save_load_checkpoint.py - updated imports
• test_trainer.py - rewrote test_normalize_mini_batch_size

---

Gradient scaling

Since the worker no longer knows mini batches, we scale gradients instead of loss:

Old (scale loss during backward)

for i in 1..N:
    grad += (1/N) * ∂loss_i/∂param
optimizer.step(grad)

New (scale gradients at optim_step)

for i in 1..N:
    grad += ∂loss_i/∂param
grad *= 1/N
optimizer.step(grad)

Both produce: grad = (1/N) * Σ ∂loss_i/∂param

---

What's next for training refactor?

• Remove weight sync logic from trainer. This should not be explicitly triggered by the trainer
• Remove other “infra” calls from trainer, such as empty_cache
• Create separate entry points for a) launching workers and b) launching training. Currently trainer.py sets up the workers (build_models) and launches training
• Update Megatron (@erictang000) to bring to same state instead of branching on which backend is used

[tyler-griggs]

This also needs to be moved out of trainer. The trainer shouldn't know/care about dp size.

[erictang000]

looks pretty good to me! will work on the megatron worker update off of this branch

Also calling out some work that's actively ongoing for adding dynamic batch sizing that's touching relevant codepaths:
#817
#847

maybe the best thing to do would be to rebase #847 on top of the tinkerify branch? (and move down the micro batching logic updates from there into into forward_backward)

[erictang000]

I get that this is cleaner to not worry about loss scaling during forward (which requires pre-knowing the details of how we break down the microbatch or how often we call optim_step).

But are you sure that this is exactly numerically equivalent + compatible with autocasting/mixed precision? Scaling the loss prior to backward seems cleaner in that sense, can we verify that this is numerically identical (and that there aren't overflow edge cases here) to scaling the loss in pytorch?

[erictang000]

coming back to this now that i'm implementing this for megatron - I think it overall seems cleaner and more flexible to handle scaling loss for gradient accumulation in forward_backward, where we could just set the default to weight each micro batch evenly and scale the loss for each microbatch by (1/num_microbatches). We could even theoretically provide an optional parameter for forward_backward that specifies a loss scaling term (default 1), so that users could flexibly scale the loss for different mini-batches differently (i.e. if you have 2 mini batches of different sizes, you might want to make sure that the 2nd mini batch loss is scaled to size).

def _forward_backward_micro(experience, microbatch_weight):
     ...
     loss = loss * microbatch_weight
     loss.backward()
     ...
def forward_backward(data, loss_weight=1.0):
     micro_batch_iterator = BatchIterator(data, micro_batch_size, drop_last=False)
     for micro_batch in micro_batch_iterator:
            metrics = self._forward_backward_micro(micro_batch, microbatch_weight=loss_weight / len(micro_batch_iterator))
      ...

This would still be compatible with the upstream tinker API which doesn't let you specify a loss weight.

wdyt? @tyler-griggs

[erictang000]

one reason that this logic was so messy before was that you could technically overlap the forward pass for critic/ref/policy, which is something you lose here. This probably wasn't obvious or used very much in most of our configs since everything was colocated, but could matter more at scale.

Once we make the api async (and handle some requests queue on the server side) we could add that functionality back - do you think that's worth adding a TODO?

[tyler-griggs]

Yeah I was thinking the same thing. I think it's fine to lose overlapping for a bit, especially given that critic usage is presumably quite low. But it will be possible to add it back once we update the async request handling.

[gemini-code-assist]

Code Review

This pull request introduces a significant architectural refactoring aimed at improving the separation of concerns within the training stack. Key changes include moving algorithmic logic to the trainer and infrastructure logic to workers, and introducing a WorkerDispatch class to coordinate between them. This refactoring also removes redundant training methods from workers and shifts micro-batching logic there.

Additionally, this PR adds support for expert parallelism in skyrl-tx, including a ragged_dot wrapper that handles sharded experts and updates to LoRA layers to support this.

The changes are well-structured and accompanied by new tests for the WorkerDispatch functionality and expert parallelism layers. The dependency updates for vllm, torch, and other packages are consistent with the code modifications. The addition of project-summary.md provides an excellent overview of the refactoring.

Overall, this is a high-quality refactoring that improves the architecture and adds new capabilities. I have not found any issues that require changes.