Skip to content

Upgrade post training deps - #4184

Open
khatwanimohit wants to merge 2 commits into
mainfrom
mohit/tpu-infer-07-16
Open

Upgrade post training deps#4184
khatwanimohit wants to merge 2 commits into
mainfrom
mohit/tpu-infer-07-16

Conversation

@khatwanimohit

@khatwanimohit khatwanimohit commented Jun 17, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

Summary

  • GCC version: Upgrades TPU Dockerfile base image from python:3.12-slim-bullseye
    (GCC 10.2.1) to python:3.12-slim-bookworm (GCC 12.2.0). The new vllm commit
    (a30addc7) introduced a CMake check requiring GCC ≥ 11.3 for C++20 support
    vllm PR #44923.

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

build_maxtext_docker_image WORKFLOW=post-training completes successfully

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

codecov Bot commented Jun 17, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 52.94118% with 16 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...ners/post_train/distillation/distillation_utils.py 56.66% 5 Missing and 8 partials ⚠️
.../trainers/post_train/distillation/train_distill.py 25.00% 2 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch from 3421155 to 96eb4e0 Compare June 30, 2026 03:45
@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch from 96eb4e0 to 6eb9a55 Compare July 23, 2026 17:23
@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch from 6eb9a55 to 6e45c71 Compare July 23, 2026 17:25
@khatwanimohit khatwanimohit changed the title Fix post-training Docker build for new vllm commit Upgrade post training deps Jul 23, 2026
@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch 3 times, most recently from 77d3e9f to a19e7b9 Compare July 23, 2026 19:10
@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch from a19e7b9 to 1fd6bbc Compare July 23, 2026 19:32
@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch from df8f2e2 to fd87931 Compare July 24, 2026 18:23
@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch 3 times, most recently from 0a12419 to 1f6332b Compare July 27, 2026 16:54
SurbhiJainUSC
SurbhiJainUSC reviewed Jul 27, 2026
View reviewed changes
Comment thread src/dependencies/scripts/install_post_train_extra_deps.py
current_dir = os.path.dirname(os.path.abspath(__file__))
repo_root = os.path.abspath(os.path.join(current_dir, "..", ".."))
github_deps_path = os.path.join(repo_root, "dependencies", "extra_deps", "post_train_github_deps.txt")
if not os.path.exists(github_deps_path):

@SurbhiJainUSC SurbhiJainUSC Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this check removed?

@khatwanimohit khatwanimohit Jul 27, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added the check back

NuojCheng
NuojCheng reviewed Jul 27, 2026
View reviewed changes
Comment thread tests/integration/checkpoint_conversion_test.py
"override_model_config=True",
"scan_layers=False",
"hardware=cpu",
"attention=dot_product",

@NuojCheng NuojCheng Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why this change?

@khatwanimohit khatwanimohit Jul 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this was fix the CI failure, but Hengtao has merged this to main in #4609. So I'm removing this and rebasing

@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch from 1f6332b to 255d2e1 Compare July 27, 2026 18:35
igorts-git
igorts-git reviewed Jul 27, 2026
View reviewed changes
Comment thread src/dependencies/scripts/install_post_train_extra_deps.py Outdated

is_root = os.geteuid() == 0 if hasattr(os, "geteuid") else False
has_sudo = shutil.which("sudo") is not None
if (is_root or has_sudo) and shutil.which("apt-get"):

@igorts-git igorts-git Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would rather not run sudo commands automatically on behalf of our users. Maybe just print a error message with instructions?
The logic below is very specific for running on Linux systems based on Debian 12 and 11. I am not sure how it would work on other distributions or on future versions of Debian.

@khatwanimohit khatwanimohit Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removed sudo usage and moved the installation command to docker image

@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch 2 times, most recently from 3c1a7f5 to 136b878 Compare July 29, 2026 01:02
@khatwanimohit
khatwanimohit force-pushed the mohit/tpu-infer-07-16 branch from 136b878 to 328b59e Compare July 29, 2026 16:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@SurbhiJainUSC SurbhiJainUSC SurbhiJainUSC left review comments
@igorts-git igorts-git igorts-git left review comments
@NuojCheng NuojCheng NuojCheng left review comments
@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner
@parambole parambole Awaiting requested review from parambole parambole is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@darisoy darisoy Awaiting requested review from darisoy darisoy is a code owner
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner
@xibinliu xibinliu Awaiting requested review from xibinliu xibinliu is a code owner
@huytransformer huytransformer Awaiting requested review from huytransformer huytransformer is a code owner
@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner
@gagika gagika Awaiting requested review from gagika gagika is a code owner
@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner
@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner
@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@suexu1025 suexu1025 Awaiting requested review from suexu1025
@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 is a code owner
@shuningjin shuningjin Awaiting requested review from shuningjin

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@khatwanimohit @SurbhiJainUSC @igorts-git @NuojCheng