Add Trainer Protocol

[allenwang28]

This PR introduces the Trainer Protocol in src/forge/api/trainer.py, establishing a unified training interface that all trainer implementations will conform to.

Motivation

Currently, Forge applications directly use Monarch actors (e.g., RLTrainer.options(...).as_actor(...)), which exposes implementation details like .route() and .fanout() to application code.

This creates tight coupling and makes it difficult to:

• Switch between different trainer backends (TorchTitan, HuggingFace, etc.)
• Write portable application code that doesn't depend on Monarch specifics

Protocol + Wrappers

Note that we're using Python's Protocol, and not ABC! In case you weren't aware, there is a big philosophical debate about ABC vs Protocol that Claude has introduced me to. I'm primarily choosing Protocol because it's lighter weight (and let me know if you disagree strongly).

Why Protocol and not ABC?

We want a public, lightweight interface that anyone can satisfy without importing our base class. Protocol gives us structural typing: if it quacks, it flies. An ABC would force nominal inheritance and encourage a hierarchy we don’t actually need.

TL;DR:

• Looser coupling: Call sites accept “anything with the right methods,” not “anything that inherits our base.”
• Frictionless third-party impls: External teams can implement the interface without depending on our internals.
• Small, composable capabilities: Easy to define narrow traits and mix them.
• Optional runtime checks: If desired, @runtime_checkable enables isinstance(x, Trainer) as a light guard.

What it looks like in practice:

With ABC:

 # Would force inheritance
  class TitanTrainer(Trainer):  # Must inherit from ABC
      def __init__(self, actor_handle):
          super().__init__()  # ABC initialization overhead
          self._actor = actor_handle

With Protocol:

  # No inheritance required
  class TitanTrainer:  # Just a plain class
      def __init__(self, actor_handle):
          self._actor = actor_handle  # That's it

Why this matters:

• Simple/thin wrappers: HFTrainer, TitanTrainer, etc. can be simple adapters over the Monarch actor
• Fungibility by default: Third parties can drop in their own trainer without subclassing anything.
• Stability for callers: Callers type against the behavior (the protocol), so internal refactors don’t cascade.
• Escape hatch: If we later need shared behavior, we can add an optional BaseTrainer(ABC) with helpers/metrics—without changing the public Trainer protocol.
• Ultimately this all allows us to keep looser coupling between the protocol definition and implementation.

Other planned changes:

The Protocol is step 1 of a multi-PR refactor. Other planned changes:

• Restructure actor/trainer.py into trainer/titan.py and rename RLTrainer to TitanTrainerActor. Add TitanTrainer wrapper class that hides Monarch adverbs
• Implement the rest of the API for titan trainer (we only do train_step right now)
• App migration - maybe after the other API changes have landed

[joecummings]

Overall very happy with this but I need to see more details in the docstrings bc its not entirely clear whats being passed into each method.

[casteryh]

I like protocol more than ABC

[casteryh]

what's they key of dict[str, Any]?

[allenwang28]

a placeholder now is returning the path it was written to and the step. we might not need anything though i'm not sure, but I want to return a placeholder at least in case of like "hosted service" territory

[felipemello1]

if we know what the dict returns, lets make it a dataclass?

[felipemello1]

why would we need these 2?

[felipemello1]

Maybe it should be List[Metric], but i am not sure how exactly this will be used. I do think that [str, float] is too constraining.

torchforge/src/forge/observability/metrics.py

Line 85 in bb57589

class Metric:

[felipemello1]

i dont fully get this

step -> should be a global variable
learning_rate -> makes sense because we run scheduler.step(), but we can always do scheduler.get_lr
accumulated_steps -> unclear to me how it would know this info. Isnt this the number of microbatches?

[felipemello1]

these two should probably be dataclasses