Repository navigation
recipes(397b v7x): trainer perf flags (SparseCore offload thresholds, decoder-layer-input placement) and MaxText gather-once pin - #2895
Conversation
TRAINER_SC_OFFLOAD_THRESHOLDS=true (on in mlperf_gbs1024_submission.sh) keeps small all-reduces (<64 MiB) and all-gathers (<256 MiB) on the TensorCore path so they skip the busy SparseCore queue. RL trainer bench: 106.5 -> 105.1 s/step. (cherry picked from commit 6d589fb3d9c9ad9aedca29eea995f071e40bf488)
…her-once PRs - REMAT_DECODER_LAYER_INPUT (offload by default, as before) sets the trainer's decoder_layer_input; device keeps each layer's input in HBM and drops the per-layer host round trip. - MaxText pin: atwigg/mlperf + AI-Hypercomputer/maxtext#5718 + #5740-#5742, which gather the small FSDP-sharded kernels (router, query, shared expert, GDN in_proj_ba, K/V) once per scanned layer stack instead of in every layer.
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces performance optimizations for the 397B v7x trainer recipes. The changes focus on fine-tuning collective offloading thresholds, providing more control over memory placement for decoder layer inputs, and updating core dependencies to leverage recent improvements in MaxText for more efficient kernel execution. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![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 updates the MaxText requirements and introduces new performance tuning configurations in the MLPerf recipes, specifically adding options for SparseCore offload thresholds (TRAINER_SC_OFFLOAD_THRESHOLDS) and decoder layer input rematerialization (REMAT_DECODER_LAYER_INPUT). The review feedback recommends validating these environment variables to ensure they only accept valid values and fail loud and fast on invalid inputs, as well as checking for duplicate flags when appending to TRAINER_LIBTPU_INIT_ARGS.
| if [[ "${TRAINER_SC_OFFLOAD_THRESHOLDS:-false}" == "true" ]]; then | ||
| export TRAINER_LIBTPU_INIT_ARGS="${TRAINER_LIBTPU_INIT_ARGS} --xla_tpu_sparse_core_all_reduce_offload_min_size_in_bytes=67108864 --xla_tpu_sparse_core_all_gather_offload_min_size_in_bytes=268435456" | ||
| fi |
There was a problem hiding this comment.
Validate the TRAINER_SC_OFFLOAD_THRESHOLDS environment variable to ensure it is strictly set to true or false (failing loud and fast on invalid inputs), and check if the corresponding SparseCore flags are already present in TRAINER_LIBTPU_INIT_ARGS before appending them to avoid duplicates.
| if [[ "${TRAINER_SC_OFFLOAD_THRESHOLDS:-false}" == "true" ]]; then | |
| export TRAINER_LIBTPU_INIT_ARGS="${TRAINER_LIBTPU_INIT_ARGS} --xla_tpu_sparse_core_all_reduce_offload_min_size_in_bytes=67108864 --xla_tpu_sparse_core_all_gather_offload_min_size_in_bytes=268435456" | |
| fi | |
| case "${TRAINER_SC_OFFLOAD_THRESHOLDS:-false}" in | |
| true) | |
| if [[ " ${TRAINER_LIBTPU_INIT_ARGS} " != *" --xla_tpu_sparse_core_all_reduce_offload_min_size_in_bytes="* ]]; then | |
| export TRAINER_LIBTPU_INIT_ARGS="${TRAINER_LIBTPU_INIT_ARGS} --xla_tpu_sparse_core_all_reduce_offload_min_size_in_bytes=67108864" | |
| fi | |
| if [[ " ${TRAINER_LIBTPU_INIT_ARGS} " != *" --xla_tpu_sparse_core_all_gather_offload_min_size_in_bytes="* ]]; then | |
| export TRAINER_LIBTPU_INIT_ARGS="${TRAINER_LIBTPU_INIT_ARGS} --xla_tpu_sparse_core_all_gather_offload_min_size_in_bytes=268435456" | |
| fi | |
| ;; | |
| false) | |
| ;; | |
| *) | |
| echo "ERROR: TRAINER_SC_OFFLOAD_THRESHOLDS must be true or false, got '${TRAINER_SC_OFFLOAD_THRESHOLDS}'" >&2 | |
| exit 1 | |
| ;; | |
| esac |
References
- Fail loud and fast: Never mask missing attributes, invalid types, or broken invariants with silent default values. Let Python raise ... or raise an explicit ValueError / specific exception immediately. (link)
| # (decoder_layer_input). offload (default) parks it in host memory; device keeps it in HBM, which | ||
| # drops the per-layer host round trip (~115 ms of DMA wait per micro-batch on the RL trainer bench, | ||
| # AI-Hypercomputer/maxtext#5701) for ~1 GiB of HBM (AOT, tpu7x-256). | ||
| export REMAT_DECODER_LAYER_INPUT="${REMAT_DECODER_LAYER_INPUT:-offload}" |
There was a problem hiding this comment.
Validate the REMAT_DECODER_LAYER_INPUT environment variable to ensure it is strictly set to either offload or device to fail loud and fast on invalid inputs.
| export REMAT_DECODER_LAYER_INPUT="${REMAT_DECODER_LAYER_INPUT:-offload}" | |
| export REMAT_DECODER_LAYER_INPUT="${REMAT_DECODER_LAYER_INPUT:-offload}" | |
| case "${REMAT_DECODER_LAYER_INPUT}" in | |
| offload|device) | |
| ;; | |
| *) | |
| echo "ERROR: REMAT_DECODER_LAYER_INPUT must be 'offload' or 'device', got '${REMAT_DECODER_LAYER_INPUT}'" >&2 | |
| exit 1 | |
| ;; | |
| esac |
References
- Fail loud and fast: Never mask missing attributes, invalid types, or broken invariants with silent default values. Let Python raise ... or raise an explicit ValueError / specific exception immediately. (link)
Stacked on #2878 (
parambole/ours275-trainer-config); retarget toatwigg/mlperfonce it lands.Trainer performance changes for the 397B v7x recipes, measured or compiled on the RL trainer bench (AI-Hypercomputer/maxtext#5701, 128 chips):
TRAINER_SC_OFFLOAD_THRESHOLDSmlperf_397b_256_v7x.sh, on inmlperf_gbs1024_submission.shdecoder_layer_input)REMAT_DECODER_LAYER_INPUT(offload/device)offload(unchanged)MoE token chunking (a candidate lever for the dependency-bound MoE token all-gather/reduce-scatter) needs no change here:
NUM_MOE_TOKEN_CHUNKS=4already overrides the tuned value.Dry run (
DRY_RUN=true ... mlperf_gbs1024_submission.sh startwithREMAT_DECODER_LAYER_INPUT=device): both threshold flags appear in the trainer env, the Pathways workers and the proxy;decoder_layer_input=devicein the trainer's MaxText flags.Bench runs of the combined changes are pending capacity.