Add maestro framework

[x41lakazam]

Maestro is a Pytorch-based benchmarking framework for overlapping communication/compute operations in distributed environments.

This presents a more accurate way to measure performance of AI patterns compared to the micro-benchmarking which measures performance of standalone operations.
It allows a user to define a workload pattern, then benchmark performance of every running element.

Example use-case of an overlapping benchmarking suite

As an example, one user would want to see how well AllGather and ReduceScatter overlap given a certain fabric and algorithms, maestro allow him to easily define it as a pattern running the two and measure the performance of each while they are overlapping, he would get such a result:

Block            | Size (B) | Avg lat (ms) | Min lat | Max lat | P99 lat | Avg BW (GB/s)
allgather        | 4M       | 1.4          | 0.2     | 3.0     | 3.0     | 10.8
reduce_scatter   | 4M       | 1.5          | 0.3     | 3.1     | 3.1     | 6.0

And therefore be able to see the impact of tuning each part of the system on the collectives, and optimize it.

Quick description of maestro capabilities

Maestro benchmark patterns, which are composed of blocks, i.e reusable unit of works that enqueue a GPU activity, and axes, i.e a collection of GPU teams, the user defines which blocks should run on which axis, maestro run the full pattern and report the performance of each block, separately.

An axis is representing the way of splitting the ranks in a classic parallelized AI workload, for example:

Who benefits from maestro

	Communication Engr.	AI researcher	Marketing teams
Goal	Maximize collective performance on realistic end-to-end workloads	Identify the best overlap opportunities (compute/communication and comm./comm.)	Quantify and communicate some added value on realistic overlap patterns
Means	Analyze collective behavior under complex overlap schedules	Benchmark experimental overlap patterns and report performance/efficiency	Compare “classic” patterns across ecosystems (ex: different algorithms) and summarize the impact
Stack used	Pre-configured kernels; customize the YAML file	Custom experimental kernels	No development - uses YAML provided by AI experts
Knowledge required	Communication libraries internals	AI kernels and execution/overlap behavior	None (ability to interpret benchmark results)

Use cases

Below is a few use cases that leveraged maestro for optimization:

• Overlapping SHARP ReduceScatter + AllGather: SHARP team used maestro to better tune both algorithm and reach a full 2 x SOL bandwidth, classical benchmark didn't offer the possibility to run both collectives at the same time properly
• QoS Priorization of collectives: In a highly parallelized AI workload, priorization is key, as such collectives are blocking entire layers while some others are less critical, maestro helped correctly leveraging QoS for priorizing some collectives over other ones
• Congestion control algorithms: Overlapping communications introduced new congestion patterns that didn't manifest in the standalone execution of basic collectives, maestro helped identifying the pain points and optimizing CC algorithms

[meta-cla]

Hi @x41lakazam!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[x41lakazam]

Hey @jansel @williamwen42 @bobrenjc93, is there any missing step to move this PR to review ?

[kwen2501]

Hi @d4l3k @dcci @bobrenjc93
This is a communication-compute overlap benchmark from NV.
Would it help evaluate collectives towards realistic settings?
Please guide us if there is a better place to contribute it. Thanks!

[lappazos]

@d4l3k @dcci @bobrenjc93 could you please review?

[d4l3k]

Generally this seems reasonable to me, have some open questions I'm discussing offline with @kwen2501 about maintenance/ownership of this including megatron pieces

[d4l3k]

Note

This review was generated by an AI agent (Claude) using PyTorch's pr-review skill against head b1ab47c, and posted at the request of @d4l3k. Findings were checked against the source but not run on GPUs; please verify before acting on them.

PR Review: #2702

Summary

Adds userbenchmark/maestro, a YAML-driven benchmark that runs collectives/GEMMs concurrently on separate process groups and reports per-block latency/BW from CUPTI. No REVIEW.md applies. Several measurement-correctness bugs mean the reported numbers are wrong (warmup included, GEMM shapes halved, start gate no-op, aliased buffers), the TorchBench entrypoint as documented doesn't run, and rank-local errors can hang the job.

Code Quality

• core/orchestrator.py:118-124, core/profilers/base.py:92: profiler.reset() only clears self._results; CUPTI buffers (10MB, cupti.py:40) are only delivered on activity_flush_all in _exit. Warmup activities are therefore still buffered at reset and come back with the measured ones. _measure_blocks_performance sorts by start_ns and pop(0)s, so it consumes the warmup iterations first; with warmup_iters >= iters the report is 100% warmup. Leftover activities are dropped silently. Fix: cupti.activity_flush_all(1) before clearing in reset() (or separate profiler contexts), and assert every mapped stream's activity list is empty after consumption.
• core/orchestrator.py:396-400: axes wait_event(start_event) before start_event is ever recorded (recorded at :427, after enqueue). torch.cuda.Event is lazily created and CUDAEvent::block() is a no-op when !is_created_, so nothing gates the launch; blocks start in CPU enqueue order and inter-axis start skew biases overlap/shared-BW. Gate on the GPU before enqueue (e.g. torch.cuda._sleep on a gate stream, record, then have axes wait).
• core/config_parser.py:134,138: GEMM dims go through parse_size(dim) // dtype_size, so mat_a_shape: [4096, 4096] becomes 2048x2048 (8x fewer FLOPs). Dimensions are element counts; don't divide. ops_presets/moe.py:56,67,75 emits element counts that are then halved the same way, and ignores its declared k/expert_capacity_factor.
• core/block.py:197-198, 246-247, core/orchestrator.py:438-456: every block slices the single main pool from offset 0. Within a block src/dst overlap (all_gather_into_tensor src=dst[:n/team] is only NCCL's valid in-place layout on team rank 0; reduce_scatter_tensor likewise; all_to_all_single/send-recv ring use identical src/dst). Across blocks, concurrent blocks on different axes and GEMM operands read/write the same memory simultaneously. Size the pool as the sum of per-block regions and give each block disjoint src/dst (optionally an explicit, correctly-offset in-place mode).
• core/orchestrator.py:263-268, 312-333: per block, start_ns is min over ranks and duration_ns max over ranks, then overlap uses start + duration — start and duration come from different ranks, and CUPTI timestamps are per-host clocks not comparable across nodes. Compute overlap per rank from that rank's own start/end, then aggregate.
• core/orchestrator.py:303-304: shared BW divides by max(duration_i) rather than the first-start-to-last-end span, so fully serialized blocks still report ~sum of individual BWs.
• core/orchestrator.py:205-217, 362-392: block↔activity mapping assumes each block launches exactly the same number of activities per stream as in one isolated run; any variation (cuBLAS split-K, NCCL memcpy-based paths) silently misattributes activities to other blocks/iterations. Only underflow is detected.
• core/orchestrator.py:145-148: latency_precision (documented in docs/doc.md:61, plumbed through main.py) is never used; output is hardcoded :.1f ms, so anything under 50µs prints 0.0.
• core/config_parser.py:31-47, core/axis.py:65: uneven teams (e.g. [[0,1,2],[3]]) are accepted, but team_size() returns len(groups[0]), giving wrong buffer sizes/BW factors on other teams. Require equal team sizes or use len(my_group_ranks).
• blocks/torch/comm.py TorchReduceScatter/TorchAllToAll: no divisibility check on vector_size vs team size (TorchAllGather has one); failure surfaces as a raw torch shape error at first run.
• core/block.py:101: _enrich_api_params_schema is never called, so api_params_schema/defaults are dead and missing keys raise bare KeyError from _prepare_tensors_pool, outside main.py's config error handling. The "Validate - all the axes referenced in the blocks exist" comment at config_parser.py:140 has no code.
• core/utils/format.py parse_size: accepts negatives (-4M → [:-N] slicing), rejects 1.5M/1MiB, and uses 1024-based units while the table header says "Size (B)" and BW uses 1e9.
• core/orchestrator.py:438, core/block.py:195,242: dtype hardcoded fp16 (fallback path allocates fp32); pool initialized with torch.ones, so repeated in-place fp16 sum reductions overflow to inf within the warmup, and (due to aliasing) GEMM operands become inf/NaN.
• core/profilers/torch.py: TorchProfiler.__init__ raises NotImplementedError and main.py only accepts cupti; remove it.
• core/axis.py:30-31: _temporary_environ clears and rebuilds all of os.environ while process-group background threads may call getenv; save/restore only the keys in envs.
• core/axis.py:53-60: a rank not in any team leaves my_group_ranks unset → AttributeError instead of the intended error.
• core/utils/distributed.py:184, 220-226: traffic_class handling is unreachable (create_axis is called without params).
• core/orchestrator.py:94: docstring says non-root ranks return None; every rank returns the full DataFrame, making run.py:25 dead.

Infrastructure

• run.py:15,18, README.md:11: python run_benchmark.py maestro-benchmark <config.yaml> fails — run_benchmark.py dispatches on the directory name (maestro). There's also no mention that it must be launched under torchrun/srun.
• maestro/maestro/__init__.py:8-10: core/, core/utils/, blocks/torch/ lack __init__.py and internal modules import top-level core, blocks, main. A namespace package loses to any regular core/blocks module anywhere on sys.path regardless of insertion order, and the permanent sys.path insert leaks these generic names to the whole process. Additionally setup.py's find_packages(include=["torchbenchmark*","userbenchmark*"]) skips the namespace dirs, so a non-editable pip install . of TorchBench ships an unimportable maestro. Use package-relative imports, add the __init__.py files, drop the sys.path hack.
• run.py:28,37: metrics JSON puts everything under a new top-level results key with an empty metrics dict; ADDING_USERBENCHMARKS.md and TorchBench consumers (e.g. abtest.py) only read metrics: {name: float}. Flatten to e.g. f"{pattern}-{block}-avg_latency_ms".
• No install.py/requirements entry for click, cupti-python; pyproject.toml is not used on the TorchBench path, and run_benchmark.py's except ImportError turns the resulting failures into exit 0.
• deploy/Dockerfile:1,7,13: git clone <REPO> is a shell redirect and fails; cloning pytorch/benchmark would pip install TorchBench's setup.py, not maestro; empty PYTORCH_VERSION default yields an invalid image ref; PIP_CONSTRAINT= drops NGC constraints for all deps. Use COPY and install userbenchmark/maestro.
• examples/execution_scripts/slurm.sbatch:82: runs /PATH/TO/maestro/src/main.py (no src/ exists) and bypasses TorchBench; --results-dir is validated but unused; SLURM_GPUS_PER_NODE is unset/typed under --gres; $CONFIG_FILE unquoted inside bash -c.

Examples

• examples/config_files/collectives_1D/*.yaml: save_results_to: "/path/to/results/..." fails at parse time with PermissionError (not in main.py's caught set); all_reduce_sweep.yaml/reduce_sweep.yaml require 16 ranks without saying so; reduce_sweep.yaml iters: 3 makes P99 meaningless; ag_rs.yaml is named "AG+RS (256M)" but its all-gather is 4M.

Backward Compatibility

• pyproject.toml:15: torch>=2.6, but Work.block_current_stream (used every _execute, core/utils/async_work.py:22) only exists from 2.9, and backend.mem_allocator (register_buffers) from 2.7. On <2.9 every pattern raises AttributeError and is skipped. Require torch>=2.9.
• pyproject.toml:12,14: exact pins click==8.3.0, pandas==2.3.3 will force downgrades/upgrades in a TorchBench env; use lower bounds.

Thread Safety

• core/orchestrator.py:81-86 (stop_on_error defaults false, config_parser.py:226): an exception on a subset of ranks (per-rank CUPTI capture failures at :376-383 / :198-215, OOM, the uneven-team case above) makes that rank skip to the next pattern and enter the world barrier/all_gather_object while peers are still in axis collectives — collective mismatch → hang until the NCCL watchdog or corrupt results. Make per-rank checks reach a common status collective and raise on all ranks together; don't continue after a block has issued collectives.
• core/utils/distributed.py:15-21: SLURM_PROCID/SLURM_NTASKS take precedence over RANK/WORLD_SIZE, so srun ... torchrun gives every local worker the node index as rank → duplicate ranks/hang. Prefer RANK when set. :147 also selects the device by global_rank % device_count; use LOCAL_RANK/SLURM_LOCALID.
• blocks/torch/comm.py:95: batch_isend_irecv works are discarded, so axis.stream never depends on the P2P kernels; the per-iteration wait_stream + barrier fence (orchestrator.py:415-424) doesn't cover the send-recv ring and iterations drift across barriers. Wait the works inside use_axis().
• core/axis.py:18-31: per-axis envs are applied only around new_group. Most NCCL_* knobs are NCCL_PARAMs cached process-wide on first read (during world init), so per-axis values are ignored, and use_axis never applies them despite the orchestrator.py:48 docstring. Document/validate which vars can vary per axis or use ProcessGroupNCCL.Options().config via pg_options.
• main.py:139-152, run.py:38, config_parser.py:14: every rank writes save_results_to and dump_output's second-resolution metrics-<ts>.json to the same paths; %D is computed per rank at import with minute resolution. Write from rank 0 only and broadcast the timestamp.

API Design

• main.py:95-101: the exported run_benchmark calls sys.exit(1) on config errors and logs a python main.py run -c usage line; raise instead.
• main.py:145-147: to_json(path, index=False) with default orient writes the columns/index-keyed shape (different from run.py's orient="records"); to_excel needs openpyxl, which isn't a dependency, and fails after the benchmark completes.
• maestro/__init__.py:12-21: version parsed from ../pyproject.toml is 0.0.0+unknown after a wheel install (the Dockerfile path). Use importlib.metadata.
• Docs: docs/doc.md/developer.md reference src/... paths, python main.py, uv sync/uv.lock, src/blocks/torch/__init__.py — none exist; main.py docstrings use list_blocks (click exposes list-blocks); README title still "NCCL Multi Test"; doc.md:194 says preset name is a suffix (it's a prefix).

Recommendation

Request Changes

The reported metrics are currently wrong (warmup contamination, halved GEMMs, non-functional start gate, aliased buffers), the documented TorchBench invocation and non-editable install don't work, and rank-local errors can deadlock. Separately, the head commit ("Sync maestro docs/pyproject from nccl-multi-test") indicates the code is mirrored from an external upstream — the ownership/sync model should be settled before landing.

[x41lakazam]

@d4l3k addressed.

[x41lakazam]

@d4l3k The CI fail look unrelated to maestro

[lappazos]

@d4l3k Can we merge?