Repository navigation
[Recipes][GSM8K] Plumb attention DP knob and MoE optimization switches - #2822
wenxindongwork wants to merge 1 commit into
Conversation
Add ROLLOUT_ATTENTION_DP (default true) and MoE optimization switches (TPU_SMALL_ALL_REDUCE, TPU_MOE_SINGLE_SORT, TPU_MOE_ROUTER_TOPK, TPU_MOE_LOWTOK) to the GSM8K 35B recipe and launcher. - ROLLOUT_ATTENTION_DP=false updates rollout sharding to disable DP attention, raises max_num_seqs to 64, and appends kv_tp_size=4 to MAXTEXT_EXTRA_FLAGS. - MoE switches default to 0 / sort to preserve baseline behavior. - Propagates optimization flags and VLLM_RAY_EXTRA_ENV_VARS_TO_COPY to rollout containers and Ray workers. Signed-off-by: wenxindongwork <wenxindong@google.com>
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 enhances the GSM8K distributed training recipe and its associated Kubernetes launcher by adding configurable knobs for attention data parallelism and MoE-specific optimizations. These changes improve flexibility for different model configurations and ensure that performance-tuning flags are correctly passed to the rollout containers and distributed workers, maintaining consistency with existing infrastructure patterns. 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 introduces configuration options for attention data parallelism and propagates MoE and collective optimization environment variables to Ray workers in the Kubernetes launcher and recipe scripts. Feedback focuses on adhering to the 'One well-lit path' style guide rule by avoiding speculative dictionary lookups and defaults in Python snippets, simplifying the Bash loop used to build the Ray environment variable list, and removing redundant environment variable assignments in ROLLOUT_EXTRA_ENV since they are already explicitly forwarded by the launcher.
| VLLM_ADDITIONAL_CONFIG="$(python3 -c ' | ||
| import json, os | ||
| cfg = json.loads(os.environ.get("VLLM_ADDITIONAL_CONFIG") or "{}") | ||
| strategy = cfg.setdefault("sharding", {}).setdefault("sharding_strategy", {}) | ||
| strategy["enable_dp_attention"] = False | ||
| strategy["tensor_parallelism"] = 1 | ||
| print(json.dumps(cfg, separators=(",", ":"))) | ||
| ')" |
There was a problem hiding this comment.
To adhere to the One well-lit path guideline in the style guide, avoid using speculative .setdefault() with fabricated defaults (like empty dictionaries {}) on structured schemas. Instead, explicitly initialize the nested keys if they are missing.
| VLLM_ADDITIONAL_CONFIG="$(python3 -c ' | |
| import json, os | |
| cfg = json.loads(os.environ.get("VLLM_ADDITIONAL_CONFIG") or "{}") | |
| strategy = cfg.setdefault("sharding", {}).setdefault("sharding_strategy", {}) | |
| strategy["enable_dp_attention"] = False | |
| strategy["tensor_parallelism"] = 1 | |
| print(json.dumps(cfg, separators=(",", ":"))) | |
| ')" | |
| VLLM_ADDITIONAL_CONFIG="$(python3 -c ' | |
| import json, os | |
| cfg = json.loads(os.environ.get("VLLM_ADDITIONAL_CONFIG") or "{}") | |
| if "sharding" not in cfg: | |
| cfg["sharding"] = {} | |
| if "sharding_strategy" not in cfg["sharding"]: | |
| cfg["sharding"]["sharding_strategy"] = {} | |
| strategy = cfg["sharding"]["sharding_strategy"] | |
| strategy["enable_dp_attention"] = False | |
| strategy["tensor_parallelism"] = 1 | |
| print(json.dumps(cfg, separators=(",", ":"))) | |
| ')" |
References
- One well-lit path (No defensive fallbacks): Never use speculative getattr(obj, 'field', default), dict.get('key', default) on structured schemas, dict(getattr(obj, 'metadata', None) or {}), .setdefault(...) with fabricated defaults, or isinstance(..., Mapping) duck-typing branches. Access attributes and keys directly on a single, strictly typed contract. (link)
| _rollout_kv_tp="$(python3 -c ' | ||
| import json, os | ||
| s = json.loads(os.environ["VLLM_ADDITIONAL_CONFIG"])["sharding"]["sharding_strategy"] | ||
| print(int(s.get("expert_parallelism", 1)) * int(s["tensor_parallelism"])) | ||
| ')" |
There was a problem hiding this comment.
To adhere to the One well-lit path guideline in the style guide, avoid speculative dict.get('key', default) on structured schemas. Instead, explicitly check for the key or handle it cleanly.
| _rollout_kv_tp="$(python3 -c ' | |
| import json, os | |
| s = json.loads(os.environ["VLLM_ADDITIONAL_CONFIG"])["sharding"]["sharding_strategy"] | |
| print(int(s.get("expert_parallelism", 1)) * int(s["tensor_parallelism"])) | |
| ')" | |
| _rollout_kv_tp="$(python3 -c ' | |
| import json, os | |
| s = json.loads(os.environ["VLLM_ADDITIONAL_CONFIG"])["sharding"]["sharding_strategy"] | |
| ep = s["expert_parallelism"] if "expert_parallelism" in s else 1 | |
| tp = s["tensor_parallelism"] | |
| print(int(ep) * int(tp)) | |
| ')" |
References
- One well-lit path (No defensive fallbacks): Never use speculative getattr(obj, 'field', default), dict.get('key', default) on structured schemas, dict(getattr(obj, 'metadata', None) or {}), .setdefault(...) with fabricated defaults, or isinstance(..., Mapping) duck-typing branches. Access attributes and keys directly on a single, strictly typed contract. (link)
| for _opt_var in TPU_SMALL_ALL_REDUCE TPU_MOE_SINGLE_SORT TPU_MOE_ROUTER_TOPK TPU_MOE_LOWTOK; do | ||
| if [[ -n "${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY:-}" ]]; then | ||
| export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY},${_opt_var}" | ||
| else | ||
| export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="${_opt_var}" | ||
| fi | ||
| done |
There was a problem hiding this comment.
The loop propagating optimization flags to VLLM_RAY_EXTRA_ENV_VARS_TO_COPY can be simplified to a single, more readable if-else block. This avoids executing a loop in bash and improves maintainability.
| for _opt_var in TPU_SMALL_ALL_REDUCE TPU_MOE_SINGLE_SORT TPU_MOE_ROUTER_TOPK TPU_MOE_LOWTOK; do | |
| if [[ -n "${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY:-}" ]]; then | |
| export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY},${_opt_var}" | |
| else | |
| export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="${_opt_var}" | |
| fi | |
| done | |
| if [[ -n "${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY:-}" ]]; then | |
| export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY},TPU_SMALL_ALL_REDUCE,TPU_MOE_SINGLE_SORT,TPU_MOE_ROUTER_TOPK,TPU_MOE_LOWTOK" | |
| else | |
| export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="TPU_SMALL_ALL_REDUCE,TPU_MOE_SINGLE_SORT,TPU_MOE_ROUTER_TOPK,TPU_MOE_LOWTOK" | |
| fi |
| ROLLOUT_RAIDEN_ENV="USE_RAIDEN_FFI=false RAIDEN_USE_FFI=0" | ||
| fi | ||
| export ROLLOUT_EXTRA_ENV="${ROLLOUT_EXTRA_ENV:-} PYTHONUNBUFFERED=1 TUNIX_IS_INTERNAL_ENV=false USE_BATCHED_RPA_LONG_CTX_KERNEL=${USE_BATCHED_RPA_LONG_CTX_KERNEL} RAGGED_GATHER_REDUCE_VERSION=${RAGGED_GATHER_REDUCE_VERSION} RAIDEN_DEVICES_PER_HOST=${RAIDEN_DEVICES_PER_HOST} PREFUSE_MOE_WEIGHTS=${ROLLOUT_PREFUSE_MOE_WEIGHTS} ROLLOUT_PREFUSE_MOE_WEIGHTS=${ROLLOUT_PREFUSE_MOE_WEIGHTS} ROLLOUT_MESH_TP=${ROLLOUT_MESH_TP} ROLLOUT_TENSOR_PARALLEL_SIZE=${ROLLOUT_MESH_TP} ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING} VLLM_MAX_NUM_SEQS=${VLLM_MAX_NUM_SEQS} VLLM_GPU_MEMORY_UTILIZATION=${VLLM_GPU_MEMORY_UTILIZATION} NUM_PRECOMPILE_WORKERS=${NUM_PRECOMPILE_WORKERS} NEW_MODEL_DESIGN=${NEW_MODEL_DESIGN} ATTN_BUCKETIZED_NUM_REQS=${ATTN_BUCKETIZED_NUM_REQS} ATTN_CUSTOM_NUM_REQS_BUCKETS=${ATTN_CUSTOM_NUM_REQS_BUCKETS} ONEHOT_MOE_PERMUTE_THRESHOLD=${ONEHOT_MOE_PERMUTE_THRESHOLD} VLLM_MOE_CHUNK_SIZE=${VLLM_MOE_CHUNK_SIZE} SLICE_ROPE_CACHE=${SLICE_ROPE_CACHE} DP_SCHED_BATCH_PREFILL=${DP_SCHED_BATCH_PREFILL} VLLM_ENABLE_V1_MULTIPROCESSING=${VLLM_ENABLE_V1_MULTIPROCESSING} VLLM_LOGGING_LEVEL=${VLLM_LOGGING_LEVEL} LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\" ${ROLLOUT_RAIDEN_ENV}" | ||
| export ROLLOUT_EXTRA_ENV="${ROLLOUT_EXTRA_ENV:-} PYTHONUNBUFFERED=1 TUNIX_IS_INTERNAL_ENV=false USE_BATCHED_RPA_LONG_CTX_KERNEL=${USE_BATCHED_RPA_LONG_CTX_KERNEL} RAGGED_GATHER_REDUCE_VERSION=${RAGGED_GATHER_REDUCE_VERSION} RAIDEN_DEVICES_PER_HOST=${RAIDEN_DEVICES_PER_HOST} PREFUSE_MOE_WEIGHTS=${ROLLOUT_PREFUSE_MOE_WEIGHTS} ROLLOUT_PREFUSE_MOE_WEIGHTS=${ROLLOUT_PREFUSE_MOE_WEIGHTS} ROLLOUT_MESH_TP=${ROLLOUT_MESH_TP} ROLLOUT_TENSOR_PARALLEL_SIZE=${ROLLOUT_MESH_TP} ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING} VLLM_MAX_NUM_SEQS=${VLLM_MAX_NUM_SEQS} VLLM_GPU_MEMORY_UTILIZATION=${VLLM_GPU_MEMORY_UTILIZATION} NUM_PRECOMPILE_WORKERS=${NUM_PRECOMPILE_WORKERS} NEW_MODEL_DESIGN=${NEW_MODEL_DESIGN} ATTN_BUCKETIZED_NUM_REQS=${ATTN_BUCKETIZED_NUM_REQS} ATTN_CUSTOM_NUM_REQS_BUCKETS=${ATTN_CUSTOM_NUM_REQS_BUCKETS} ONEHOT_MOE_PERMUTE_THRESHOLD=${ONEHOT_MOE_PERMUTE_THRESHOLD} VLLM_MOE_CHUNK_SIZE=${VLLM_MOE_CHUNK_SIZE} SLICE_ROPE_CACHE=${SLICE_ROPE_CACHE} DP_SCHED_BATCH_PREFILL=${DP_SCHED_BATCH_PREFILL} VLLM_ENABLE_V1_MULTIPROCESSING=${VLLM_ENABLE_V1_MULTIPROCESSING} VLLM_LOGGING_LEVEL=${VLLM_LOGGING_LEVEL} LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\"${TPU_SMALL_ALL_REDUCE:+ TPU_SMALL_ALL_REDUCE=${TPU_SMALL_ALL_REDUCE}}${TPU_MOE_SINGLE_SORT:+ TPU_MOE_SINGLE_SORT=${TPU_MOE_SINGLE_SORT}}${TPU_MOE_ROUTER_TOPK:+ TPU_MOE_ROUTER_TOPK=${TPU_MOE_ROUTER_TOPK}}${TPU_MOE_LOWTOK:+ TPU_MOE_LOWTOK=${TPU_MOE_LOWTOK}}${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY:+ VLLM_RAY_EXTRA_ENV_VARS_TO_COPY=\"${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY}\"} ${ROLLOUT_RAIDEN_ENV}" |
There was a problem hiding this comment.
Since k8s_launcher.sh has been updated to explicitly plumb and forward TPU_SMALL_ALL_REDUCE, TPU_MOE_SINGLE_SORT, TPU_MOE_ROUTER_TOPK, TPU_MOE_LOWTOK, and VLLM_RAY_EXTRA_ENV_VARS_TO_COPY (lines 620-625 in k8s_launcher.sh), appending them to ROLLOUT_EXTRA_ENV here is completely redundant. Removing them simplifies this extremely long line and avoids double-passing these variables.
| export ROLLOUT_EXTRA_ENV="${ROLLOUT_EXTRA_ENV:-} PYTHONUNBUFFERED=1 TUNIX_IS_INTERNAL_ENV=false USE_BATCHED_RPA_LONG_CTX_KERNEL=${USE_BATCHED_RPA_LONG_CTX_KERNEL} RAGGED_GATHER_REDUCE_VERSION=${RAGGED_GATHER_REDUCE_VERSION} RAIDEN_DEVICES_PER_HOST=${RAIDEN_DEVICES_PER_HOST} PREFUSE_MOE_WEIGHTS=${ROLLOUT_PREFUSE_MOE_WEIGHTS} ROLLOUT_PREFUSE_MOE_WEIGHTS=${ROLLOUT_PREFUSE_MOE_WEIGHTS} ROLLOUT_MESH_TP=${ROLLOUT_MESH_TP} ROLLOUT_TENSOR_PARALLEL_SIZE=${ROLLOUT_MESH_TP} ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING} VLLM_MAX_NUM_SEQS=${VLLM_MAX_NUM_SEQS} VLLM_GPU_MEMORY_UTILIZATION=${VLLM_GPU_MEMORY_UTILIZATION} NUM_PRECOMPILE_WORKERS=${NUM_PRECOMPILE_WORKERS} NEW_MODEL_DESIGN=${NEW_MODEL_DESIGN} ATTN_BUCKETIZED_NUM_REQS=${ATTN_BUCKETIZED_NUM_REQS} ATTN_CUSTOM_NUM_REQS_BUCKETS=${ATTN_CUSTOM_NUM_REQS_BUCKETS} ONEHOT_MOE_PERMUTE_THRESHOLD=${ONEHOT_MOE_PERMUTE_THRESHOLD} VLLM_MOE_CHUNK_SIZE=${VLLM_MOE_CHUNK_SIZE} SLICE_ROPE_CACHE=${SLICE_ROPE_CACHE} DP_SCHED_BATCH_PREFILL=${DP_SCHED_BATCH_PREFILL} VLLM_ENABLE_V1_MULTIPROCESSING=${VLLM_ENABLE_V1_MULTIPROCESSING} VLLM_LOGGING_LEVEL=${VLLM_LOGGING_LEVEL} LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\"${TPU_SMALL_ALL_REDUCE:+ TPU_SMALL_ALL_REDUCE=${TPU_SMALL_ALL_REDUCE}}${TPU_MOE_SINGLE_SORT:+ TPU_MOE_SINGLE_SORT=${TPU_MOE_SINGLE_SORT}}${TPU_MOE_ROUTER_TOPK:+ TPU_MOE_ROUTER_TOPK=${TPU_MOE_ROUTER_TOPK}}${TPU_MOE_LOWTOK:+ TPU_MOE_LOWTOK=${TPU_MOE_LOWTOK}}${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY:+ VLLM_RAY_EXTRA_ENV_VARS_TO_COPY=\"${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY}\"} ${ROLLOUT_RAIDEN_ENV}" | |
| export ROLLOUT_EXTRA_ENV="${ROLLOUT_EXTRA_ENV:-} PYTHONUNBUFFERED=1 TUNIX_IS_INTERNAL_ENV=false USE_BATCHED_RPA_LONG_CTX_KERNEL=${USE_BATCHED_RPA_LONG_CTX_KERNEL} RAGGED_GATHER_REDUCE_VERSION=${RAGGED_GATHER_REDUCE_VERSION} RAIDEN_DEVICES_PER_HOST=${RAIDEN_DEVICES_PER_HOST} PREFUSE_MOE_WEIGHTS=${ROLLOUT_PREFUSE_MOE_WEIGHTS} ROLLOUT_PREFUSE_MOE_WEIGHTS=${ROLLOUT_PREFUSE_MOE_WEIGHTS} ROLLOUT_MESH_TP=${ROLLOUT_MESH_TP} ROLLOUT_TENSOR_PARALLEL_SIZE=${ROLLOUT_MESH_TP} ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING} VLLM_MAX_NUM_SEQS=${VLLM_MAX_NUM_SEQS} VLLM_GPU_MEMORY_UTILIZATION=${VLLM_GPU_MEMORY_UTILIZATION} NUM_PRECOMPILE_WORKERS=${NUM_PRECOMPILE_WORKERS} NEW_MODEL_DESIGN=${NEW_MODEL_DESIGN} ATTN_BUCKETIZED_NUM_REQS=${ATTN_BUCKETIZED_NUM_REQS} ATTN_CUSTOM_NUM_REQS_BUCKETS=${ATTN_CUSTOM_NUM_REQS_BUCKETS} ONEHOT_MOE_PERMUTE_THRESHOLD=${ONEHOT_MOE_PERMUTE_THRESHOLD} VLLM_MOE_CHUNK_SIZE=${VLLM_MOE_CHUNK_SIZE} SLICE_ROPE_CACHE=${SLICE_ROPE_CACHE} DP_SCHED_BATCH_PREFILL=${DP_SCHED_BATCH_PREFILL} VLLM_ENABLE_V1_MULTIPROCESSING=${VLLM_ENABLE_V1_MULTIPROCESSING} VLLM_LOGGING_LEVEL=${VLLM_LOGGING_LEVEL} LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\" ${ROLLOUT_RAIDEN_ENV}" |
This branch has not been deployed
2 participants
Summary
Plumb
ROLLOUT_ATTENTION_DPknob and MoE optimization switches to the GSM8K distributed recipe (gsm_8k_35b_256.sh) and math launcher (math_gsm8k_dist/k8s_launcher.sh), mirroring the plumbing added tomlperf_base.shanddeepswe_dist/k8s_launcher.shin PRs #2806 and #2811.ROLLOUT_ATTENTION_DP=falseconfigures rolloutsharding_strategywithenable_dp_attention: false,tensor_parallelism: 1, andexpert_parallelism: 4.max_num_seqs: 64.kv_tp_size=4intoMAXTEXT_EXTRA_FLAGSfor the MaxText trainer to align KV head shapes during weight sync.TPU_SMALL_ALL_REDUCE,TPU_MOE_SINGLE_SORT,TPU_MOE_ROUTER_TOPK, andTPU_MOE_LOWTOK(defaulting to0/"sort").VLLM_RAY_EXTRA_ENV_VARS_TO_COPYthroughmath_gsm8k_dist/k8s_launcher.shto the rollout containers and vLLM Ray TPU workers.Verification
bash -non both scripts.DRY_RUN=true) under default parameters and withROLLOUT_ATTENTION_DP=false ROLLOUT_FP8=true TRAINER_FP8=true TPU_SMALL_ALL_REDUCE=1 TPU_MOE_SINGLE_SORT=1 TPU_MOE_ROUTER_TOPK=rowmax TPU_MOE_LOWTOK=1 MAX_STEPS=10.