Implement forward-only pass and populate metrics

[tyler-griggs]

Summary

• Implement forward() in SkyRL-Train backend (was raising NotImplementedError)
• forward_backward() now returns training metrics (total_loss:sum, pg_loss:sum, entropy_loss:sum, num_tokens:sum) extracted from the dispatch worker return dict
• optim_step() now returns grad_norm and learning_rate in its metrics field
• Added metrics field to OptimStepOutput type

[pcmoritz]

This is an important fix, don't forget to revert the change before merging the PR

[tyler-griggs]

Oh thank you, I was just hacking around a bit with Claude, it will be reverted

[gemini-code-assist]

Code Review

This pull request implements the forward pass in the SkyRL-Train backend and adds metric reporting to forward_backward and optim_step. The changes are well-structured and align with the goals outlined in the description. I've suggested a minor refactoring in _extract_metrics to improve maintainability by using a dictionary for metric mapping, which will make it easier to add new metrics in the future. Overall, this is a solid contribution.

[gemini-code-assist]

The current implementation of _extract_metrics uses a series of if statements to map metric names. While this works, it can become less maintainable as the number of metrics grows. Using a dictionary to define the mapping from source metric names to target metric names would make the code more scalable and easier to read.

    def _extract_metrics(self, data: dict) -> dict[str, float]:
        """Extract training metrics from dispatch return dict.

        Workers return metrics like 'loss', 'policy_loss', 'policy_entropy', etc.
        We convert to Tinker's colon-suffixed format (e.g. 'total_loss:sum').
        """
        metrics: dict[str, float] = {}
        metric_mapping = {
            "policy_loss": "pg_loss:sum",
            "policy_entropy": "entropy_loss:sum",
            "response_length": "num_tokens:sum",
        }

        # SFT path returns 'loss'; RL path returns 'final_loss' / 'policy_loss'
        if "loss" in data:
            metrics["total_loss:sum"] = float(data["loss"])
        elif "final_loss" in data:
            metrics["total_loss:sum"] = float(data["final_loss"])

        for source_key, target_key in metric_mapping.items():
            if source_key in data:
                metrics[target_key] = float(data[source_key])

        return metrics