perf(rollout): add HybridEngine rollout profiling (#8295)
## Summary
This PR adds opt-in stage-level profiling for `HybridEngineRollout`.
The profiling path measures rollout-level latency without changing the
default execution behavior. It establishes a measurable baseline for the
HybridEngine rollout investigations discussed in #8197.
Following maintainer feedback, the executable OPSD HybridEngine rollout
benchmark has been moved to DeepSpeedExamples and is no longer part of
this PR.
Companion benchmark PR:
<[DEEPSPEED_EXAMPLES_PR_URL](https://github.com/deepspeedai/DeepSpeedExamples/pull/1009)>
## Motivation
OPSD-style workloads commonly generate multiple responses for each
prompt. Before optimizing this path, we need a reproducible way to
measure the rollout stages and determine where time is spent.
The initial profiling API records:
- prompt batch expansion
- model generation
- rollout post-processing
- end-to-end rollout latency
- generated-token throughput
- rollout workload metadata
The executable benchmark that exercises this API across different
prompt, response, batch-size, and sample-count combinations is
maintained separately in DeepSpeedExamples.
## Changes
### Opt-in rollout profiling
This PR adds `enable_profiling` to `HybridEngineRolloutConfig`.
Profiling is disabled by default:
```python
rollout = HybridEngineRollout(engine, tokenizer)
```
It can be enabled explicitly with:
```python
config = HybridEngineRolloutConfig(enable_profiling=True)
rollout = HybridEngineRollout(engine, tokenizer, config)
```
When enabled, `HybridEngineRollout` records synchronized measurements
for:
- `prompt_expansion_ms`
- `generation_ms`
- `post_processing_ms`
- `total_ms`
- `tokens_per_second`
The profile also records:
- input batch size
- samples per prompt
- prompt length
- returned response length
- total generated-token count
Profiling remains disabled by default because accelerator
synchronization affects normal execution performance.
The most recent measurement can be retrieved with:
```python
profile = rollout.get_last_profile()
```
When profiling is disabled, the normal rollout execution path and output
behavior remain unchanged.
### Correctness
The rollout now preserves a tokenizer `pad_token_id` of `0` instead of
treating it as missing and replacing it with the EOS token.
Tests cover:
- profiling disabled by default
- profiling enabled and disabled paths
- output equivalence with profiling enabled
- synchronized timing fields
- multiple samples per prompt
- generated-token counts
- prompt and attention-mask alignment
- zero-valued pad token IDs
- `get_last_profile()` behavior
## Companion benchmark
The executable OPSD HybridEngine rollout benchmark has been moved to
DeepSpeedExamples following maintainer feedback:
<[DEEPSPEED_EXAMPLES_PR_URL](https://github.com/deepspeedai/DeepSpeedExamples/pull/1009)>
The companion benchmark supports configurable matrices for:
- batch size
- samples per prompt
- prompt length
- response length
- FP16 or BF16
- warmup iterations
- measured iterations
- inference-cache retention or release
It reports:
- prompt expansion latency
- generation latency
- post-processing latency
- total rollout latency
- generated-token throughput
- peak accelerator memory
- raw per-iteration profiles
- mean, p50, and p95 summaries
The benchmark executes the largest effective batch first so HybridEngine
initializes a sufficiently large inference workspace, while preserving
the user-requested order in the output JSON.
Its initial validation scope is intentionally limited to:
- one accelerator process
- one GPU
- ZeRO stage 0
- exact-length synthetic prompts
The benchmark depends on the profiling API introduced by this PR.
## Validation
Test environment:
- Python 3.12.3
- Pytest 9.1.1
- Transformers 4.40.2
Command:
```bash
pytest -q tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
```
Result:
```text
15 passed
```
The modified files also pass the repository pre-commit hooks, including:
- YAPF
- clang-format
- flake8
- codespell
- license checks
- torch distributed import checks
- accelerator abstraction checks
The executable benchmark and its CPU-only tests are validated separately
in the companion DeepSpeedExamples PR.
## Scope
This PR introduces only opt-in rollout-level profiling and its DeepSpeed
core correctness coverage.
It does not include an executable benchmark in the DeepSpeed core
repository.
It does not attempt to optimize generation or attribute time to internal
HybridEngine operations such as:
- parameter gathering
- LoRA fuse/unfuse transitions
- inference-cache acquisition or release
- prefill
- decode
- CUDA graph execution
Those internal phases can be investigated separately after the profiling
API and companion benchmark establish a reproducible baseline.
This PR does not modify the existing inference-cache lifecycle or claim
to fix cache release and reacquisition behavior.
Related to #8197.
---------
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
