perf(opsd): add HybridEngine rollout profiling benchmark by nathon-lee · Pull Request #1009 · deepspeedai/DeepSpeedExamples

nathon-lee · 2026-08-23T12:52:41Z

Summary

This PR adds a standalone OPSD HybridEngine rollout profiling benchmark to DeepSpeedExamples.

The executable benchmark was moved out of the DeepSpeed core repository following maintainer feedback on DeepSpeed PR #8295. DeepSpeed PR #8295 retains the opt-in HybridEngineRollout profiling API, while this PR provides the benchmark runner, documentation, and CPU-only tests.

Related to:

Dependency

This benchmark depends on the rollout profiling API introduced by DeepSpeed PR #8295:

deepspeedai/DeepSpeed#8295

Until that PR is merged, the benchmark must be run with its DeepSpeed checkout first on PYTHONPATH.

Changes

Add benchmarks/opsd/benchmark_hybrid_engine_rollout.py
Add benchmark usage and scope documentation
Add CPU-only tests for:
- argument parsing
- percentile summaries
- largest-effective-batch execution ordering
- JSON result fields

The benchmark supports matrices for:

batch size
samples per prompt
prompt length
response length
FP16 and BF16
warmup and 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 iteration profiles
mean, p50, and p95 summaries

The largest effective batch executes first to initialize the HybridEngine workspace, while JSON results preserve the order requested on the command line.

Scope

The initial benchmark scope is intentionally limited to:

one process
one GPU
ZeRO stage 0
synthetic exact-length prompts

This measures HybridEngine rollout-level performance. It is not a complete OPSD training-step benchmark and does not measure teacher inference, loss computation, backward, or optimizer work.

This PR does not modify DeepSpeed core code or redesign the rollout profiling API.

Validation

Static and CPU checks

python benchmarks/opsd/benchmark_hybrid_engine_rollout.py --help

python -m py_compile \
  benchmarks/opsd/benchmark_hybrid_engine_rollout.py \
  benchmarks/opsd/tests/test_benchmark_hybrid_engine_rollout.py

python -m unittest discover \
  -s benchmarks/opsd/tests \
  -p 'test_*.py' \
  -v

git diff --check

Result:

Ran 4 tests
OK

Single-GPU smoke test

Environment:

Python 3.12
PyTorch 2.9.1+cu128
Transformers 4.40.2
Accelerate 1.14.0
NVIDIA RTX A6000
one process
ZeRO stage 0

Command:

PYTHONPATH=/workspace/DeepSpeed_woo:/workspace/DeepSpeedExamples_woo \
python -m torch.distributed.run \
  --nproc_per_node=1 \
  benchmarks/opsd/benchmark_hybrid_engine_rollout.py \
  --model facebook/opt-125m \
  --batch-sizes 1 \
  --samples-per-prompt 1 4 \
  --prompt-lengths 8 \
  --response-lengths 1 \
  --warmup 1 \
  --iterations 1 \
  --output /tmp/opsd_examples_samples_matrix.json

Result:

benchmark completed successfully
JSON output contained both requested cases
samples_per_prompt=1 generated 1 token
samples_per_prompt=4 generated 4 tokens
raw profiles and mean/p50/p95 summaries were written
requested result order was preserved

Validation boundaries

The FP16 retained-cache path was validated on one RTX A6000.

BF16 and larger model/sequence matrices were not run in this validation environment.

The --release-inference-cache option is forwarded to the existing DeepSpeed HybridEngine configuration. In this environment, the current DeepSpeed PR branch failed in the existing cache-retake path with a RecursionError before the measured iteration. This PR does not modify or claim to fix that DeepSpeed core behavior.

pre-commit was unavailable in the original development environment. The available static checks and CPU tests listed above passed.

cc @delock

Signed-off-by: nathon-lee <leejianwoo@gmail.com>

## 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 deepspeedai#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](deepspeedai/DeepSpeedExamples#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](deepspeedai/DeepSpeedExamples#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 deepspeedai#8197. --------- Signed-off-by: nathon-lee <leejianwoo@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

delock · 2026-08-24T10:49:19Z

Hi @nathon-lee , thank you for moving the benchmark code here. I leave some comments to followup. Thanks!

Signed-off-by: nathon-lee <leejianwoo@gmail.com>

nathon-lee · 2026-08-24T11:45:59Z

Thanks for the review, @delock. I've removed p95 from the throughput summary and dropped the tests, leaving only the benchmark and README.

perf(opsd): add HybridEngine rollout profiling benchmark

45c74e6

Signed-off-by: nathon-lee <leejianwoo@gmail.com>

nathon-lee requested a review from tjruwase as a code owner August 23, 2026 12:52

nathon-lee mentioned this pull request Aug 23, 2026

perf(rollout): add HybridEngine rollout profiling deepspeedai/DeepSpeed#8295

Merged

delock reviewed Aug 24, 2026

View reviewed changes

Comment thread benchmarks/opsd/benchmark_hybrid_engine_rollout.py Outdated

delock reviewed Aug 24, 2026

View reviewed changes

Comment thread benchmarks/opsd/tests/test_benchmark_hybrid_engine_rollout.py Outdated

delock requested changes Aug 24, 2026

View reviewed changes

perf(opsd): address benchmark review feedback

efdb5ab

Signed-off-by: nathon-lee <leejianwoo@gmail.com>

Navigation Menu

Uh oh!

perf(opsd): add HybridEngine rollout profiling benchmark - #1009

perf(opsd): add HybridEngine rollout profiling benchmark#1009
nathon-lee wants to merge 2 commits into
deepspeedai:masterfrom
nathon-lee:perf/opsd-hybridengine-rollout-benchmark

nathon-lee commented Aug 23, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

delock commented Aug 24, 2026

Uh oh!

nathon-lee commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

nathon-lee commented Aug 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Dependency

Changes

Scope

Validation

Static and CPU checks

Single-GPU smoke test

Validation boundaries

Uh oh!

Uh oh!

Uh oh!

delock commented Aug 24, 2026

Uh oh!

nathon-lee commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

nathon-lee commented Aug 23, 2026 •

edited

Loading