diff --git a/benchmarks/opsd/README.md b/benchmarks/opsd/README.md index e58be3800..c683885ba 100644 --- a/benchmarks/opsd/README.md +++ b/benchmarks/opsd/README.md @@ -2,10 +2,11 @@ `benchmark_hybrid_engine_rollout.py` measures rollout-level performance for an OPSD workload backed by DeepSpeed HybridEngine. It runs a matrix of synthetic, -exact-length prompts and reports prompt expansion, generation, post-processing, -total latency, generated-token throughput, and peak accelerator memory. Each -case includes raw iteration profiles, mean and p50 summaries, and p95 summaries -for latency metrics. +exact-length prompts and reports prompt expansion, generation, prefill/decode +forward timings, generation overhead, post-processing, total latency, +generated-token throughput, and peak accelerator memory. Each case includes +raw iteration profiles, mean/p50/p95 summaries for latency metrics, and a +consistent `num_decode_forwards` count. This benchmark depends on the rollout profiling API introduced by DeepSpeed PR #8295: @@ -13,7 +14,8 @@ DeepSpeed PR #8295: https://github.com/deepspeedai/DeepSpeed/pull/8295 Use a DeepSpeed checkout that contains that API and place it first on -`PYTHONPATH`. The current validation scope is one process, one GPU, and ZeRO-0. +`PYTHONPATH`. The current validation scope is one process, one GPU, and ZeRO-0 +with `inference_tp_size=1`. This is a HybridEngine rollout benchmark, not a complete OPSD training-step benchmark; it does not measure teacher inference, loss computation, backward, or optimizer work. @@ -27,18 +29,89 @@ control unreported warmup calls and recorded calls. Pass `--release-inference-cache` to release the inference cache after generation; otherwise the benchmark retains it. `--temperature`, `--top-p`, `--seed`, and `--output` control sampling, reproducibility, and the JSON output path. +Pass `--use-shared-prefill` to enable shared prompt prefill. Shared-prefill +currently supports only ZeRO-0, `inference_tp_size=1`, and the internal KV +cache; it cannot be combined with `--release-inference-cache`. The selected +mode is recorded as `use_shared_prefill` in the top-level JSON result. +Pass `--use-graph-capture` to benchmark CUDA Graph decode. CUDA Graph capture +can be combined with `--use-shared-prefill`, but supports only greedy, +fixed-length generation (`--temperature 0`). The first generation performs +graph capture, so use `--warmup 1` or greater before recording measurements. The largest effective batch (`batch_size * samples_per_prompt`) executes first so HybridEngine initializes a sufficiently large inference workspace. Results in the output JSON retain the matrix order requested on the command line. -From the DeepSpeedExamples repository root, run a single-GPU benchmark with: +## Profile fields + +The raw profile for every iteration contains these latency fields in +milliseconds: + +- `prompt_expansion_ms`: time to expand prompts for `n_samples_per_prompt`. +- `generation_ms`: time spent in generation after prompt expansion. +- `prefill_forward_ms`: the first top-level model forward (the prefill). +- `decode_forward_ms`: cumulative time of subsequent top-level model forwards. +- `generation_overhead_ms`: `generation_ms` minus prefill and decode forward + time, covering the remaining generation work. +- `post_processing_ms`: time to construct the rollout result after generation. +- `total_ms`: end-to-end time for the profiled rollout. + +`num_decode_forwards` is metadata counting the forwards included in +`decode_forward_ms`; it is retained in each raw profile and summarized once per +case after verifying that every iteration has the same count. It is not a +latency statistic. Forward breakdown fields can be `null` on paths such as +CUDA Graph replay. If all iterations are unavailable, their case summary is +`null`; a mixture of `null` and numeric values is rejected. + +For A/B comparisons, keep model, dtype, seed, workload matrix, warmup, and +iterations exactly identical between runs; change only the feature under test. + +From the DeepSpeedExamples repository root, run the baseline single-GPU OPT-6.7B +FP16 benchmark with: + +```bash +PYTHONPATH=/workspace/DeepSpeed_woo:/workspace/DeepSpeedExamples \ +torchrun --nproc_per_node=1 \ + benchmarks/opsd/benchmark_hybrid_engine_rollout.py \ + --model facebook/opt-6.7b \ + --dtype fp16 \ + --batch-sizes 1 \ + --samples-per-prompt 1 4 \ + --prompt-lengths 128 512 \ + --response-lengths 32 128 \ + --temperature 0 \ + --warmup 5 \ + --iterations 20 \ + --output /workspace/results/opsd_prefill_decode_baseline.json +``` + +Run the shared-prefill variant with the same workload and controls: + +```bash +PYTHONPATH=/workspace/DeepSpeed_woo:/workspace/DeepSpeedExamples \ +torchrun --nproc_per_node=1 \ + benchmarks/opsd/benchmark_hybrid_engine_rollout.py \ + --model facebook/opt-6.7b \ + --dtype fp16 \ + --batch-sizes 1 \ + --samples-per-prompt 1 4 \ + --prompt-lengths 128 512 \ + --response-lengths 32 128 \ + --temperature 0 \ + --warmup 5 \ + --iterations 20 \ + --use-shared-prefill \ + --output /workspace/results/opsd_prefill_decode_shared.json +``` + +For a quick local smoke test, reduce warmup and iterations: ```bash PYTHONPATH=/workspace/DeepSpeed_woo:/workspace/DeepSpeedExamples \ torchrun --nproc_per_node=1 \ benchmarks/opsd/benchmark_hybrid_engine_rollout.py \ --model facebook/opt-6.7b \ + --dtype fp16 \ --batch-sizes 1 \ --samples-per-prompt 1 4 \ --prompt-lengths 128 512 \ diff --git a/benchmarks/opsd/benchmark_hybrid_engine_rollout.py b/benchmarks/opsd/benchmark_hybrid_engine_rollout.py index 3d4b610ff..b8e17aa82 100644 --- a/benchmarks/opsd/benchmark_hybrid_engine_rollout.py +++ b/benchmarks/opsd/benchmark_hybrid_engine_rollout.py @@ -17,8 +17,16 @@ from pathlib import Path -_LATENCY_FIELDS = ("prompt_expansion_ms", "generation_ms", "post_processing_ms", "total_ms") -_THROUGHPUT_FIELDS = ("tokens_per_second", ) +_LATENCY_FIELDS = ( + "prompt_expansion_ms", + "generation_ms", + "prefill_forward_ms", + "decode_forward_ms", + "generation_overhead_ms", + "post_processing_ms", + "total_ms", +) +_THROUGHPUT_FIELDS = ("tokens_per_second",) def _percentile(values, percentile): @@ -31,6 +39,11 @@ def _summarize(profiles): summary = {} for field in _LATENCY_FIELDS: values = [profile[field] for profile in profiles] + if all(value is None for value in values): + summary[field] = None + continue + if any(value is None for value in values): + raise ValueError(f"Profile field {field} is None for only some iterations") summary[field] = { "mean": statistics.mean(values), "p50": statistics.median(values), @@ -45,6 +58,13 @@ def _summarize(profiles): return summary +def _validate_decode_forward_counts(profiles): + counts = [profile["num_decode_forwards"] for profile in profiles] + if len(set(counts)) != 1: + raise ValueError(f"num_decode_forwards differs across iterations: {counts}") + return counts[0] + + def _build_parser(): parser = argparse.ArgumentParser(description="Benchmark HybridEngine rollout stage profiling") parser.add_argument("--model", default="facebook/opt-6.7b", help="HuggingFace model ID or local model path") @@ -58,6 +78,8 @@ def _build_parser(): parser.add_argument("--warmup", type=int, default=5) parser.add_argument("--iterations", type=int, default=20) parser.add_argument("--release-inference-cache", action="store_true") + parser.add_argument("--use-shared-prefill", action="store_true") + parser.add_argument("--use-graph-capture", action="store_true") parser.add_argument("--seed", type=int, default=1234) parser.add_argument("--output", default="opsd_rollout_profile.json") return parser @@ -72,6 +94,10 @@ def _validate_args(args): raise ValueError("temperature must be non-negative") if not 0.0 < args.top_p <= 1.0: raise ValueError("top-p must be in the interval (0, 1]") + if getattr(args, "use_shared_prefill", False) and args.release_inference_cache: + raise ValueError("--use-shared-prefill cannot be combined with --release-inference-cache") + if args.use_graph_capture and args.temperature != 0.0: + raise ValueError("--use-graph-capture requires --temperature 0.0 for greedy generation") def _load_model_and_tokenizer(model_name, dtype, device): @@ -106,6 +132,7 @@ def _build_engine(model, args): }, "hybrid_engine": { "enabled": True, + "enable_cuda_graph": args.use_graph_capture, "max_out_tokens": max(args.prompt_lengths) + max(args.response_lengths), "release_inference_cache": args.release_inference_cache, }, @@ -157,6 +184,7 @@ def _run_case(rollout, model, args, batch_size, samples_per_prompt, prompt_lengt "requested_response_length": response_length, "returned_response_length": profiles[-1]["response_length"], "peak_memory_mb": accelerator.max_memory_allocated() / (1024**2), + "num_decode_forwards": _validate_decode_forward_counts(profiles), "summary": _summarize(profiles), "profiles": profiles, } @@ -181,6 +209,8 @@ def _build_result(args, device, cases): "temperature": args.temperature, "top_p": args.top_p, "release_inference_cache": args.release_inference_cache, + "use_shared_prefill": getattr(args, "use_shared_prefill", False), + "use_graph_capture": args.use_graph_capture, "cases": cases, } @@ -206,7 +236,12 @@ def _run(args): try: model, tokenizer = _load_model_and_tokenizer(args.model, dtype, device) engine = _build_engine(model, args) - rollout = HybridEngineRollout(engine, tokenizer, HybridEngineRolloutConfig(enable_profiling=True)) + rollout = HybridEngineRollout( + engine, + tokenizer, + HybridEngineRolloutConfig(enable_profiling=True, + use_shared_prefill=args.use_shared_prefill), + ) case_specs, execution_order = _ordered_case_specs(args) cases = [None] * len(case_specs)