Skip to content

Fix duplicate XLA compilation of jit(train_step) #4934

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
olupton wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix duplicate XLA compilation of jit(train_step) #4934

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions src/maxtext/trainers/pre_train/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -958,7 +958,10 @@ def train_loop(config, recorder, state=None):
params_shardings,
)

with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
# Do not enter the legacy `mesh` context manager here: the training loop calls
# p_train_step without it, and the mismatch in jit's tracing-cache key would
# cause train_step to be traced and compiled a second time on the first step.
with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):

@NuojCheng NuojCheng Aug 18, 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.

@khatwanimohit is the diloco implementation still blocked by this?

data_sharding = sharding.get_input_data_sharding(config, mesh)
shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding)
if config.shard_optimizer_over_data and isinstance(model, nn.Module):
Expand All @@ -972,7 +975,12 @@ def train_loop(config, recorder, state=None):
else:
lower_args = (state, shaped_batch)
maxtext_utils.maybe_dump_jaxpr(config, p_train_step, lower_args)
if config.compiled_trainstep_file == "": # compile only when there is no pre-compiled file loaded
if config.compiled_trainstep_file == "" and not jax.config.jax_enable_pgle:
Comment thread
olupton marked this conversation as resolved.
# Compile only when there is no pre-compiled file loaded. With AutoPGLE, an
# ahead-of-time compiled executable can never be reused by the dispatch path
# (the active PGLE profiler is part of JAX's executable cache key), so this
# compile would only add a third full compilation on top of the profiling
# compile and the FDO recompile; skip it and its memory stats.
compiler_options = max_utils.parse_libtpu_flags_to_dict(config.compile_xla_flags)
compiled = p_train_step.lower(*lower_args).compile(compiler_options=compiler_options)
compiled_stats = compiled.memory_analysis()
Expand Down
Loading