Skip to content

Add PyTorch support to TPU Raiden H2D/D2H presubmit performance gate for GitHub Actions. - #925

Merged
copybara-service[bot] merged 1 commit into
mainfrom
test_981433715
Oct 15, 2026
Merged

copybara-service[bot] merged 1 commit into
mainfrom
test_981433715

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Sep 14, 2026 •

Copy link
Copy Markdown

Add PyTorch support to TPU Raiden H2D/D2H presubmit performance gate for GitHub Actions.

  • Decouple perf_core.py under function measure(...) supporting both framework="jax" and framework="torch".
  • Add PyTorch SPMD multi-processing transfer execution (_torch_worker_fn) with safe library preloading via torch_tpu_common_loader.
  • Update h2d_d2h_benchmark_gating.py with --framework=jax|torch|all.
  • Split the floors per framework into h2d_d2h_gating_jax_baselines.json and h2d_d2h_gating_torch_baselines.json.
  • Populate h2d_d2h_gating_torch_baselines.json with verified physical TPU v5e-8 (8-chip 2x4 mesh) baselines and statistical floors (1000 iterations).
  • Update benchmark_registry.pbtxt to define symmetrical JAX and PyTorch metrics for GitHub Actions / MLCompass.

@google-cla

google-cla Bot commented Sep 14, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@copybara-service
copybara-service Bot force-pushed the test_981433715 branch 4 times, most recently from 35ace3b to ce82cd0 Compare September 21, 2026 01:13
@copybara-service copybara-service Bot changed the title Add PyTorch support to TPU Raiden H2D/D2H presubmit performance gating. Add PyTorch support to TPU Raiden H2D/D2H presubmit performance floor measurement for GitHub Actions and internal testing. Sep 23, 2026
@copybara-service copybara-service Bot changed the title Add PyTorch support to TPU Raiden H2D/D2H presubmit performance floor measurement for GitHub Actions and internal testing. Add PyTorch support to TPU Raiden H2D/D2H presubmit performance gate for GitHub Actions. Sep 23, 2026
@copybara-service
copybara-service Bot force-pushed the test_981433715 branch 2 times, most recently from b0c5e08 to a829034 Compare September 23, 2026 02:17
@copybara-service
copybara-service Bot force-pushed the test_981433715 branch 2 times, most recently from 78ca817 to 44686e7 Compare October 15, 2026 03:32
…for GitHub Actions.

- Decouple `perf_core.py` under function `measure(...)` supporting both `framework="jax"` and `framework="torch"`.
- Add PyTorch SPMD multi-processing transfer execution (`_torch_worker_fn`) with safe library preloading via `torch_tpu_common_loader`.
- Update `h2d_d2h_benchmark_gating.py` with `--framework=jax|torch|all`.
- Split the floors per framework into `h2d_d2h_gating_jax_baselines.json` and `h2d_d2h_gating_torch_baselines.json`.
- Populate `h2d_d2h_gating_torch_baselines.json` with verified physical TPU v5e-8 (8-chip 2x4 mesh) baselines and statistical floors (1000 iterations).
- Update `benchmark_registry.pbtxt` to define symmetrical JAX and PyTorch metrics for GitHub Actions / MLCompass.

PiperOrigin-RevId: 999818997
@copybara-service
copybara-service Bot merged commit 3603668 into main Oct 15, 2026
1 check failed
@copybara-service
copybara-service Bot deleted the test_981433715 branch October 15, 2026 03:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

0 participants