Consolidate benchmark scripts under src/scripts/benchmarks#786
Conversation
rowwise_alltoall_test.py and rowwise_combine_fused_test.py are standalone torchrun benchmark scripts (argparse + main()), not pytest tests — they define no test functions, so pytest collected them and found nothing. Rename them to *_bench.py to match the existing benchmark-script convention (routed_experts_fp8_bench.py, shared_experts_dense_bench.py) so they are no longer collected while remaining runnable via torchrun. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move all benchmark scripts into a single src/scripts/benchmarks/ directory and add the MoE-v2 kernel microbenchmarks that were still missing. - Relocate the existing benches: odc (renamed benchmark_odc.py -> odc_bench.py), all_reduce_bench, routed_experts_fp8_bench, shared_experts_dense_bench, rowwise_alltoall_bench, rowwise_combine_fused_bench. - Add microbenches for grouped-mm backends, MXFP8 quantize (plain / swizzle / weighted / scale-mode), scaled_grouped_mm_q, rowwise FP8 dispatch/combine, fused restore+unpermute, and valid-prefix SwiGLU. These are standalone CLI / torchrun scripts (no pytest test functions), so this also removes them from test collection. Update in-file usage-example paths. Co-Authored-By: Tianhua Tao <tianhuatao@gmail.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
AkshitaB
changed the title
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Point the sys.path shim at src after the move
When this renamed benchmark is run from a source checkout via the torchrun path shown above, it now lives under src/scripts/benchmarks, so Path(__file__).resolve().parents[4] resolves to /workspace instead of the repo's src directory (the old src/test/nn/moe/v2 location made this depth work). In environments without an editable install, _load_combine_callables() inserts the wrong directory and the following from olmo_core.kernels.symm_mem_vdev2d ... import fails, so the benchmark is no longer runnable after the move.
Point the sys.path shim at src after the move
After moving this script to src/scripts/benchmarks, Path(__file__).resolve().parents[4] no longer points at /workspace/OLMo-core/src; it points at /workspace. For the source-tree torchrun use case that this helper supports, _load_rowwise_callables() therefore adds the wrong path and from olmo_core.kernels.symm_mem_vdev2d ... fails unless the package happens to be installed separately. Recompute this as the actual src directory for the new location.
ℹ️ 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".
The torch baseline sliced both `x[:num_elements]` and `out[:num_elements]` with the device-tensor count, so under torch.compile (fullgraph) each slice derived an independent unbacked size and dynamo could not prove they matched for the copy_ (Eq(u2, u1)). Reuse the sliced row count for the output slice instead. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Drop the bespoke SubCmd / BeakerLaunchConfig launch machinery so the all-reduce bench runs like the other benches: a plain argparse entrypoint under torchrun (or the generic Beaker launcher), with no required RUN_NAME/CLUSTER args. Payload size, trials, and warmup are now CLI flags; the byte size is derived from the tensor instead of module globals. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The rowwise all-to-all and combine benches inserted `parents[4]` onto sys.path, which was `src` at their old src/test/nn/moe/v2/ location. After the move to src/scripts/benchmarks/, `src` is `parents[2]`; `parents[4]` pointed above the repo, breaking the `from olmo_core.kernels.symm_mem_vdev2d ...` import in a non-editable-install (torchrun-from-source) run. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Optionally torch.compile both the FeedForward and SharedExperts forwards before timing, so the comparison can measure whether inductor fuses away the SharedExperts permute + strided SwiGLU + batch-1 bmm overhead that makes the eager E=1 path slower. Adds --compile and --compile-mode (default/reduce-overhead/max-autotune); the header notes the mode. Co-Authored-By: Tianhua Tao <tianhuatao@gmail.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The single-expert (shared) forward pays for a no-op permute and a batch-1 bmm that a plain mm would avoid; eager it runs ~13% slower than FeedForward (parity under torch.compile max-autotune). Leave a TODO to add an E==1 fast path later. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 participant
Collect all benchmark/profiling scripts into a single
src/scripts/benchmarks/directory so they're easy to find and clearly separated from the test suite.
Changes
src/scripts/benchmarks/:benchmark_odc.py→odc_bench.pyall_reduce_bench.py(fromsrc/scripts/train/)routed_experts_fp8_bench.py,shared_experts_dense_bench.py,rowwise_alltoall_bench.py,rowwise_combine_fused_bench.py(from
src/test/nn/moe/**)quantize (plain / swizzle / weighted / scale-mode),
scaled_grouped_mm_q,rowwise-FP8 dispatch/combine comm, fused restore+unpermute, and valid-prefix
SwiGLU.
All are standalone CLI /
torchrunscripts with no pytest test functions, sothey are no longer collected during test runs. In-file usage-example paths were
updated to match the new locations.