Retry with rbf = -1 if tokens are dropped - #5080
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a retry mechanism (retry_when_tokens_dropped) for Mixture of Experts (MoE) training, allowing the training step to be retried with a dropless buffer capacity if tokens are dropped due to ragged sort buffer overflow. The changes span configuration files, MoE routing layers, training loops, and unit tests. The review feedback highlights several critical issues: a performance regression caused by host-device synchronization on every step, a logging bug where metrics are overwritten during a retry, compatibility issues with pure NNX models due to the use of self.sow, a trivial unit test that does not actually execute the retry logic, and a JAX tracing anti-pattern when collecting intermediate metrics.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
5148c64 to
3f95504
Compare
RoutedMoE now falls back to a dropless (worst-case) ragged buffer and redoes just that layer's route+compute via jax.lax.cond when the tuned ragged_buffer_factor would otherwise drop tokens, gated behind retry_when_tokens_dropped. Also validates against num_moe_emb_chunks>0 (not supported), and surfaces overflow via a gated log line in training_loop_iteration.
3f95504 to
b074035
Compare
1 participant
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/555278394
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.