perf(rollout): add HybridEngine rollout profiling

[nathon-lee]

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>

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:

rollout = HybridEngineRollout(engine, tokenizer)

It can be enabled explicitly with:

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:

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>

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:

pytest -q tests/unit/runtime/rollout/test_hybrid_engine_rollout.py

Result:

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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 138d1bd7fb

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".