Add loss_fn parameterization to forward_backward

[tyler-griggs]

Summary

• Adds optional loss_fn and loss_fn_config parameters to forward_backward() for Tinker API compatibility
• Maps Tinker algorithm names (e.g., "ppo") to SkyRL equivalents
• Removes ppo_train() from FSDP workers - uses gradient scaling at optim_step instead
• Updates all tests to use the new unified API

Changes

1. WorkerDispatch (worker_dispatch.py):

   • Added loss_fn and loss_fn_config parameters to forward_backward()
   • Passes these parameters through to worker mesh methods

2. PolicyWorkerBase (worker.py):

   • Added convert_tinker_loss_config() static method to convert Tinker's absolute ratio bounds to SkyRL's offset format
   • Gradient scaling now happens at optim_step time based on accumulated micro batches
   • Removed separate ppo_train() path for FSDP workers

3. Tests:

   • Updated test helpers to use comprehensive parameter sets
   • Added test_convert_tinker_loss_config for Tinker config conversion
   • Updated all GPU tests to use pass_through routing and positional batch parameters

Test Plan

• ✅ CPU tests pass: test_normalize_mini_batch_size, test_convert_tinker_loss_config
• GPU tests can be run with: pytest tests/gpu/gpu_ci/test_training_step.py

Stack

• Built on top of [Tinker] Add save_weights_for_sampler() to WorkerDispatch #922 (save_weights_for_sampler)

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces parameterization for loss_fn to forward_backward for Tinker API compatibility and refactors weight synchronization into a new save_weights_for_sampler method. The changes are generally positive, simplifying the API and improving test coverage. However, I've identified a critical issue in trainer.py where the critic training path appears to be broken due to incorrect arguments being passed to _execute_training_step. Additionally, there's an inconsistency in a new test file regarding a test utility function call.

[gemini-code-assist]

There appears to be a critical issue with the arguments passed to _execute_training_step for both the critic and policy models. The function signature for _execute_training_step is (self, model: str, data: TrainingInputBatch), but it's being called with three arguments here (e.g., self._execute_training_step("critic", data, "critic")). This will result in a TypeError at runtime.

While the intent seems to be to pass a loss_fn, the implementation appears incomplete. Specifically:

1. The signature of _execute_training_step hasn't been updated to accept a third argument.
2. Even if it were updated, the critic training path would likely fail. The loss_fn would be "critic", which is not handled by PolicyWorkerBase._get_loss_fn, and CriticWorkerBase doesn't have a comparable method to handle a parameterized loss function.

To fix this, you'll need to update the signature of _execute_training_step and ensure that both policy and critic workers can correctly handle the new loss_fn parameter. For the critic, you might want to pass None as the loss_fn if it's not meant to be parameterized, and handle that case in _execute_training_step.

[gemini-code-assist]

The call to make_dummy_training_batch here and on line 190 seems inconsistent with changes in other test files. In other files like test_save_load_checkpoint.py and test_training_step.py, the batch_size argument was removed from this call (e.g., make_dummy_training_batch()).

If the signature of make_dummy_training_batch has changed, this could lead to test failures. For consistency across the test suite, please update this call to match the new pattern.

Suggested change

	dummy_batch = make_dummy_training_batch(batch_size=dp_size)
	dummy_batch = make_dummy_training_batch()