Skip to content

Use lower-bounds for tpu-post-train dependencies#3271

Merged
copybara-service[bot] merged 1 commit intomainfrom
bvandermoon-post-train-deps
Feb 27, 2026
Merged

Use lower-bounds for tpu-post-train dependencies#3271
copybara-service[bot] merged 1 commit intomainfrom
bvandermoon-post-train-deps

Conversation

@bvandermoon
Copy link
Copy Markdown
Collaborator

Description

Replace strict dependency matching with lower-bounds for the post-training dependencies. We already use these in the pre-training dependencies.

Tests

Ran this on a clean venv:

uv pip install -e .[tpu-post-train] --resolution=lowest
install_maxtext_tpu_post_train_github_deps

# We will try to move this to the `install_maxtext_tpu_post_train_github_deps` script
uv pip install src/maxtext/integration/vllm --no-deps


export HF_TOKEN=<hf_token>
export WORKLOAD=gemma3_rl_5
export OUTPUT_PATH=<gcs_bucket>

NEW_MODEL_DESIGN=1 TPU_BACKEND_TYPE=jax python3 -m src.maxtext.trainers.post_train.rl.train_rl src/maxtext/configs/post_train/rl.yml \
  model_name=gemma3-4b \
  tokenizer_path=google/gemma-3-4b-it \
  run_name=$WORKLOAD \
  base_output_directory=$OUTPUT_PATH \
  hf_access_token=$HF_TOKEN \
  batch_size=4 \
  num_batches=5 \
  scan_layers=False \
  hbm_utilization_vllm=0.4 \
  rollout_data_parallelism=2 \
  rollout_tensor_parallelism=2 \
  allow_split_physical_axes=true \
load_parameters_path=<gcs_checkpoint_path> \
vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
vllm_additional_config='{"maxtext_config": {"model_name": "gemma3-4b", "log_config": "false"}}' 2>&1 | tee  grpo_out_gemma3.txt

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@codecov
Copy link
Copy Markdown

codecov Bot commented Feb 27, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@bvandermoon bvandermoon force-pushed the bvandermoon-post-train-deps branch from 1fda43f to 4d33491 Compare February 27, 2026 21:00
@copybara-service copybara-service Bot merged commit 0955203 into main Feb 27, 2026
42 checks passed
@copybara-service copybara-service Bot deleted the bvandermoon-post-train-deps branch February 27, 2026 22:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants