Skip to content

refactor(trainers): share auxiliary-role lifecycle - #452

Open
yuetongz417 wants to merge 1 commit into
inclusionAI:mainfrom
yuetongz417:refactor/share-role-aware-trainer-lifecycle
Open

refactor(trainers): share auxiliary-role lifecycle#452
yuetongz417 wants to merge 1 commit into
inclusionAI:mainfrom
yuetongz417:refactor/share-role-aware-trainer-lifecycle

Conversation

@yuetongz417

Copy link
Copy Markdown

Summary

This PR extracts the auxiliary model-role lifecycle shared by DPO and PPO into reusable role-aware trainer classes.

The new RoleAwareTrainerMixin centralizes:

  • backend initialization and cleanup;
  • model-role initialization through ensure_roles;
  • role-based log-probability scoring;
  • validation that returned log-probabilities align with the input token rows.

RoleAwarePolicyTrainer combines this shared role handling with PolicyOnlyTrainer, allowing rollout-based algorithms to reuse the same infrastructure.

DPO now inherits the generic role-aware mixin for its frozen reference model, while PPO inherits the role-aware policy trainer for its actor, reference, critic, and reward roles. Their existing training behavior is unchanged.

This refactor also provides a reusable foundation for future teacher-scored workflows such as OPD, where a trainable student performs rollouts and a frozen teacher scores the generated tokens.

Testing

Added CPU tests covering:

  • DPO and PPO inheritance from the shared role-aware classes;
  • role initialization order;
  • backend cleanup when role initialization fails;
  • forwarding of scoring microbatch configuration;
  • validation of missing or misaligned log-probability rows.

Extract backend role initialization, cleanup, and validated log-probability scoring into reusable role-aware trainer classes. Migrate DPO and PPO without changing their algorithm behavior, providing a shared foundation for teacher-scored workflows such as OPD.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant