Fix GPU-only rowwise EP test failures#788
Conversation
Three unrelated issues surfaced running the symm-mem rowwise tests on a multi-GPU machine (these are gated by @requires_symm_mem_vdev2d, so they never run in CI): - Drop-count metric read: `_ep_no_sync_rowwise_drop_tokens_sum` is stored wrapped in a `_HiddenTensor` (kept off the autograd graph). Unwrap it with `unhide_from_torch(...)` before `.item()`. - FP8 combine-gather lease test: the rowwise-FP8 forward consumes prequantized expert weights, which the train loop/optimizer rebuilds each step. Call `block.refresh_rowwise_fp8_cache()` before each forward in the test. - fp32 rowwise parity tests: the weighted combine backward fuses the router weighting into `dispatchRowsPutWeighted`, which is bf16/fp16-only by design. Run these blocks/inputs in bf16 (the path's production dtype) and use bf16-appropriate tolerances instead of fp32. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Drop the OLMO_RUN_ROUTED_EXPERTS_PERF_TEST opt-in gate so the pad-layout latency regression guard runs whenever a GPU (and torch grouped-mm) is available, rather than silently skipping. Removes the now-unused os/pytest imports. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![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
Here are some automated review suggestions for this pull request.
Reviewed commit: 5b77e62dbe
ℹ️ 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".
| assert drop_tokens_sum.item() == 0 | ||
| # The metric is stored wrapped in a `_HiddenTensor` (kept off the autograd graph); | ||
| # unwrap before reading its value. | ||
| assert unhide_from_torch(drop_tokens_sum).item() == 0 |
There was a problem hiding this comment.
In the rowwise=True multi-GPU case, this unwrap lets the test get past the drop-count check, but the same function still builds both blocks and the input as float32. The subsequent loss_ep.backward() goes through the rowwise weighted combine path, whose dispatchRowsPutWeighted kernel only dispatches fp16/bf16, so this test still fails on the intended GPU setup with the same unsupported-Float backward that the commit fixes elsewhere. Cast the rowwise case to bf16 and relax the tolerances as in the other rowwise parity tests.
Useful? React with 👍 / 👎.
| assert drop_tokens_sum is not None | ||
| # The metric is stored wrapped in a `_HiddenTensor` (kept off the autograd graph); unwrap | ||
| # before reading its value. | ||
| assert unhide_from_torch(drop_tokens_sum).item() == 0 |
There was a problem hiding this comment.
After this metric unwrap, _run_rowwise_ep_dropless_matches_no_ep still continues into backward() with float32 blocks/input. On the rowwise path that backward uses dispatchRowsPutWeighted, which does not implement the Float specialization, so the GPU-only test remains broken once it reaches line 174. This helper needs the same bf16 cast/input and bf16 tolerances applied to the rowwise parity helpers in block_no_sync_test.py.
Useful? React with 👍 / 👎.
The `_HiddenTensor` unwrap alone let these two tests past the drop-count check, but they then run a fp32 backward through the rowwise weighted combine, which hits the same bf16/fp16-only `dispatchRowsPutWeighted`. Cast the rowwise case to bf16 with bf16 tolerances (the synced/no-EP fp32 path keeps its tight tolerances). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 participant
Fixes for MoE-v2 rowwise expert-parallel tests that fail on a multi-GPU machine. These are gated by
@requires_symm_mem_vdev2d(needs the compiled symm-mem extension), so they never run in CI and only surface on manual GPU runs.Three root causes
Drop-count metric read (
'_HiddenTensor' object has no attribute 'item')._ep_no_sync_rowwise_drop_tokens_sumis stored wrapped in a_HiddenTensor(kept off the autograd graph). The tests now unwrap it withunhide_from_torch(...)before.item()— matching the metrics helper.FP8 combine-gather lease test (
FP8 weight store 'w_up_gate' cache 'rhs' is not initialized). The rowwise-FP8 forward consumes prequantized expert weights, which the train loop / optimizer rebuilds each step. The test now callsblock.refresh_rowwise_fp8_cache()before each forward.fp32 rowwise parity tests (
"dispatchRowsPutWeighted" not implemented for 'Float'). The forward applies the router-prob weighting in torch (dtype-agnostic), but the combine backward fuses it intodispatchRowsPutWeighted, which is bf16/fp16-only by design (not a regression). These parity tests now build both blocks + inputs in bf16 (the path's production dtype) with bf16-appropriate tolerances.Also
OLMO_RUN_ROUTED_EXPERTS_PERF_TEST=1opt-in gate ontest_routed_experts_speed_vs_pad_positionsso the pad-layout latency regression guard runs whenever a GPU is available instead of silently skipping.Note
The bf16 parity tolerances (
atol=rtol=2e-2) are a starting point and may want tightening/loosening after a GPU run. The perf test asserts wall-clock latency, so its thresholds may need slack on a shared box.