[Refactor] refactor `TrainEngine` APIs

[garrett4wade]

Checklist

• This refactor maintains backward compatibility with all user-facing APIs.
• For large-scale refactors, I’ve prepared a phased implementation plan.

Motivation

The current TrainEngine exposes only high-level APIs like train_batch and eval_batch.
This implementation has the following drawbacks:

• Significant code duplication across public APIs and between different backends
• Insufficient flexibility for fine-grained microbatch control

To address these issues and align with frameworks like verl 0.6.0 and tinker,
we plan to refactor the TrainEngine implementation.

This refactoring will maintain backward compatibility, ensuring no impact on application-level users.

Proposed Refactor Plan

class TrainEngine:
    # Internal APIs
    def _forward_compute_mb(
        self, 
        mb_input: dict[str, Any], 
        loss_fn: Callable[[torch.Tensor, dict[str, Any]], torch.Tensor],
        loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
        **kwargs,
    ) -> tuple[torch.Tensor, Callable[[torch.Tensor], tuple[torch.Tensor, dict]]:
        # Update:
        # 1. For FSDP engine, run a single forward pass and compute the loss
        # 2. For Megatron engine, run forward on this pipeline stage and output hidden states & loss fn over outputs. 
        #     If on the last stage, compute the weighted loss. Check the `forward_step` inline function in MegatronEngine.
        ...

    # Edit: This method is not mandatory temparorily.
    #def _forward_backward_compute_mb(
    #    self, 
    #    mb_input: dict[str, Any], 
    #    loss_fn: Callable[[torch.Tensor, dict[str, Any]], torch.Tensor],
    #    loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
    #    mb_output: torch.Tensor | None, 
    #    mb_output_grads: torch.Tensor | None,
    #    overlapped: bool,
    #    **kwargs,
    #) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    #    # Runs forward-backward computation with optional overlap
    #    ...
    
    # Exposed APIs
    def optimizer_zero_grad(self):
        ...
    def optimizer_step(self):
        ...
    def lr_scheduler_step(self):
        ...
    def step_lr_scheduler(self):
        # Backward compatibility
        return self.lr_scheduler_step()

    @dataclass
    class ForwardBackwardOutputs:
        mb_outputs: list[torch.Tensor] | None
        losses: list[torch.Tensor] | None

    # Edit: signature
    def forward_backward_batch(
        data_iterator: Iterable[dict[str, torch.Tensor],
        loss_fn: Callable[[torch.Tensor, dict[str, Any]], torch.Tensor],
        loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
        output_post_hook: Callable[[torch.Tensor, dict[str, Any]], Any] | None = None,
        return_outputs: bool = False,
        # edit: added a new argument
        forward_only: bool = False,
    ) -> ForwardBackwardOutputs:
        # Reduces code duplication in `train_batch`, `forward`, and `eval_batch`.
        # For FSDP, composes internal APIs `_forward_compute_mb` and `loss.backward` in a for-loop
        # For megatron, passes `_forward_compute_mb` to megatron's `forward_backward_func` and runs it.
        ...

    # High-level training APIs (maintained for backward compatibility).
    def train_batch(
        self,
        input_: dict[str, Any],
        loss_fn: Callable[[torch.Tensor, dict[str, Any]], torch.Tensor],
        loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
    ) -> dict[str, float]:
        # 1. split micro batch
        # 2. call forward_backward_batch with loss_fn and loss_weight_fn
        # 3. call optimizer step
        ...
    @torch.no_grad()
    def eval_batch(
        self,
        input_: dict[str, Any],
        loss_fn: Callable[[torch.Tensor, dict[str, Any]], torch.Tensor],
        loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
    ) -> torch.Tensor | None:
        # 1. split micro batch
        # 2. call forward_backward_batch with loss_fn and loss_weight_fn and `forward_only=True`
        # 3. aggregate losses and return
        ...
    # Renamed for consistency; maintains backward compatibility via alias below.
    @torch.no_grad()
    def forward_batch(
        self,
        input_: dict[str, Any],
        output_seqlens: list[int] | None = None,
        post_hook: Callable[[torch.Tensor, dict[str, Any]], Any] | None = None,
        aggregate_fn: Callable[[list[Any]], Any] = torch.cat,
    ) -> Any | None:
        # 1. split micro batch
        # 2. call forward_backward_batch without loss_fn, with output_post_hook, `forward_only=True`, and `return_outputs=True`
        # 3. aggregate outputs according to output_seqlens and aggregate_fn
        ...
    @torch.no_grad()
    def forward(self, ...):
        # Backward compatibility alias
        return self.forward_batch(...)
    # Other APIs (save/load, etc.) remain unchanged.
    ...

Additional Context

• verl 0.6.0 engine API
• tinker API

[nuzant]

_forward_backward_compute_mb is not really useful when implementing our own pipeline schedules and we can remove it. We need a backward for micro batch and it could be backend agnostic.

[garrett4wade]

@nuzant Okay, but _forward_backward_compute_mb should be somehow related to the backend because we need to orchestrate different modules within the backend. This method can be omitted temperoarily

[aaaandychen]

Hi @garrett4wade , I've analyzed the issue and related files. I'd like to confirm if my implementation plan aligns with your expectations:

1. Refactor Base Class (TrainEngine):
   • Implement forward_backward_batch to handle micro-batch loops, gradient accumulation and refactor top-level APIs (train_batch, eval_batch, forward_batch) to invoke forward_backward_batch.
   • Add forward as an alias for forward_batch and step_lr_scheduler for lr_scheduler_step for backward compatibility.
   • Define internal interfaces: _forward_compute_mb and three exposed APIs (e.g. optimizer_step).
2. Refactor Subclasses (BaseHF, FSDP, Megatron):
   • Remove the redundant top-level APIs to leverage inheritance.
   • Implement the backend-specific _forward_compute_mb and three exposed APIs(e.g. optimizer_step).
     If this plan looks good to you, I will proceed with the changes.

[garrett4wade]

@aaaandychen Your plan looks great. We want to reuse the code in the current train_batch, eval_batch, and forward implementations. As such, we need to extract unified low-level APIs like _forward_compute_mb and forward_backward_batch. That's the target of this refactoring.

Here are some additional notes:

1. We have removed BaseHFEngine in a recent PR refactor: merge base_hf_engine with fsdp_engine for code cleanup #629 , so you only need to modify FSDPEngine and MegatronEngine.
2. I have edited the above code snippet for improving method signature and comments. You can check those for details.

[aaaandychen]

@garrett4wade Thanks for the clarification. I fully understand the requirements now and will start working on the implementation immediately.

[aaaandychen]

Hi @garrett4wade ,
I have a few questions regarding the implementation details of the issue and would appreciate your insights.

1. Handling input_ and data_iterator in forward_backward_batch My understanding is that data_iterator serves as an iterable for retrieving data in the process loop. When the input is a single batch (input_: dict[str, any]), does it imply that we need to split micro batch in high level API? Given that different models may require customized logic for splitting data, I plan to introduce an abstract method split_micro_batch. Would this design change be acceptable? Or could you advise on the recommended way to transform the input data here?

2. Parameter Propagation across Isolated Functions Due to the function isolation after refactoring, some parameters that were previously directly accessible now need to be propagated between methods. For example, in Megatron, forward_backward_func requires num_microbatches and seq_length, which are determined during the micro-batch splitting phase in the top-level API. To maintain consistency in method signatures, I currently use a wrapper around the iterator to carry these parameters. I would like to discuss if there is a more ideal or robust solution to handle this state propagation.

Best regards.