wehub-resource-sync
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""Rollout engines for on-policy generation during RL/distillation training.
|
|
|
|
Provides:
|
|
- :class:`RolloutEngine` — abstract base class
|
|
- :class:`RolloutRequest`, :class:`RolloutBatch`, :class:`SamplingConfig` — dataclasses
|
|
- :class:`HybridEngineRollout` — concrete implementation using DeepSpeed hybrid engine
|
|
- :func:`build_rollout` — factory that selects the engine from config
|
|
"""
|
|
|
|
from deepspeed.runtime.rollout.base import (
|
|
RolloutBatch,
|
|
RolloutConfig,
|
|
RolloutEngine,
|
|
RolloutRequest,
|
|
SamplingConfig,
|
|
)
|
|
from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout
|
|
|
|
__all__ = [
|
|
"HybridEngineRollout",
|
|
"RolloutBatch",
|
|
"RolloutConfig",
|
|
"RolloutEngine",
|
|
"RolloutRequest",
|
|
"SamplingConfig",
|
|
"build_rollout",
|
|
]
|
|
|
|
|
|
def build_rollout(rollout_cfg, student_engine=None, tokenizer=None, **kwargs):
|
|
"""Factory: construct the rollout engine specified by ``rollout_cfg.engine``.
|
|
|
|
Args:
|
|
rollout_cfg: :class:`RolloutConfig` (or any object with an ``engine``
|
|
attribute set to ``"hybrid_engine"``).
|
|
student_engine: DeepSpeed engine wrapping the student model.
|
|
tokenizer: HuggingFace tokenizer.
|
|
"""
|
|
engine_name = rollout_cfg.engine
|
|
if engine_name == "hybrid_engine":
|
|
if student_engine is None or tokenizer is None:
|
|
raise ValueError("hybrid_engine rollout needs both student_engine and tokenizer")
|
|
return HybridEngineRollout(engine=student_engine, tokenizer=tokenizer, cfg=rollout_cfg)
|
|
|
|
raise ValueError(f"Unknown rollout engine {engine_name!r}; choose from 'hybrid_engine'")
|