Skip to content

[Tunix] Check on the skip_jit and cache_nnx_graph params from kwargs. cache_nnx_graph can save the cpu host time. #2178

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

Merged
copybara-service merged 1 commit into from
Sep 15, 2026
+54 −32
Merged

[Tunix] Check on the skip_jit and cache_nnx_graph params from kwargs. cache_nnx_graph can save the cpu host time. #2178

Show file tree
Hide file tree
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
20 changes: 9 additions & 11 deletions tests/experimental/train/peft_trainer_v2_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -789,18 +789,16 @@ def _write_one_checkpoint(self, root, *, implicit_resume=True):
trainer = self._external_resume_trainer(
root, implicit_resume=implicit_resume
)
trainer.fwd_bwd(self.train_ds[0])
train_steps = trainer.update()
trainer.fwd_bwd(self.train_ds[0], cache_nnx_graph=False)
train_steps = trainer.update(cache_nnx_graph=False)
trainer.save_checkpoint(
metadata={
'step': train_steps,
'policy_version': train_steps,
'num_rollouts': 8,
}
)
trained_state = jax.tree.map(
jnp.copy, nnx.state(trainer.model, nnx.Param)
)
trained_state = jax.tree.map(jnp.copy, nnx.state(trainer.model, nnx.Param))
trainer.close()
return train_steps, trained_state

Expand Down Expand Up @@ -870,12 +868,10 @@ def test_restore_checkpoint_accepts_explicit_step(self):
trainer = self._external_resume_trainer(root, implicit_resume=True)
states = {}
for _ in range(2):
trainer.fwd_bwd(self.train_ds[0])
step = trainer.update()
trainer.fwd_bwd(self.train_ds[0], cache_nnx_graph=False)
step = trainer.update(cache_nnx_graph=False)
trainer.save_checkpoint(metadata={'step': step, 'marker': step})
states[step] = jax.tree.map(
jnp.copy, nnx.state(trainer.model, nnx.Param)
)
states[step] = jax.tree.map(jnp.copy, nnx.state(trainer.model, nnx.Param))
trainer.close()

rolled_back = self._external_resume_trainer(root, implicit_resume=False)
Expand All @@ -892,7 +888,9 @@ def test_restore_checkpoint_accepts_unexpected_kwargs(self):
config = peft_trainer_v2.TrainingConfig(eval_every_n_steps=1000)
model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
trainer = peft_trainer_v2.PeftTrainer(model, optax.sgd(1e-3), config)
metadata = trainer.restore_checkpoint(checkpoint_directory='/somewhere/else')
metadata = trainer.restore_checkpoint(
checkpoint_directory='/somewhere/else'
)
self.assertEqual(metadata, {'step': 0})

def test_restore_checkpoint_without_configured_directory_is_a_noop(self):
Expand Down
66 changes: 45 additions & 21 deletions tunix/experimental/train/peft_trainer_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,8 @@ class GradientAccumulator(nnx.Module):
Persistent vs. non-persistent mode (`self.persistent`):
Controlled by `allocate_grads` at initialization:
* Persistent mode (`allocate_grads=True`, `self.persistent=True`): Used
when accumulating across multiple micro-steps (`gradient_accumulation_steps
when accumulating across multiple micro-steps
(`gradient_accumulation_steps
> 1`). A parameter-sized buffer is allocated at initialization and zeroed
in-place when `reset()` is called so the buffer persists across updates.
* Non-persistent mode (`allocate_grads=False`, `self.persistent=False`):
Expand Down Expand Up @@ -321,10 +322,10 @@ def _add(acc_var, g_var):

if jax.tree_util.tree_leaves(self.grads):
jax.tree_util.tree_map(
_add,
self.grads,
grads,
is_leaf=lambda x: isinstance(x, nnx.Variable),
_add,
self.grads,
grads,
is_leaf=lambda x: isinstance(x, nnx.Variable),
)
else:
# No buffer held: either it was never allocated, or a non-persistent
Expand Down Expand Up @@ -353,10 +354,10 @@ def _scale_and_cast(v, target_dtype):
raise ValueError(
"The gradient accumulator is empty. Either get() was called without a"
" preceding add()/set(), or the gradients written by an earlier"
" executable were discarded on the way out of jit -- nnx.cached_partial"
" (cache_nnx_graph=True) freezes the bound module's graphdef, so a"
" step that changes the accumulator's pytree structure cannot hand it"
" to a later executable."
" executable were discarded on the way out of jit --"
" nnx.cached_partial (cache_nnx_graph=True) freezes the bound"
" module's graphdef, so a step that changes the accumulator's pytree"
" structure cannot hand it to a later executable."
)

if not jax.tree_util.tree_leaves(self._param_dtypes):
Expand Down Expand Up @@ -789,7 +790,7 @@ def _shard(x, p):
)

def jit_fwd_bwd_update_and_eval_step(
self, skip_jit: bool = False, cache_nnx_graph: bool = False
self, skip_jit: bool = False, cache_nnx_graph: bool = True
):
"""Creates and returns the train and eval step functions.

Expand Down Expand Up @@ -823,7 +824,8 @@ def jit_fwd_bwd_update_and_eval_step(
else:
donate_argnames = ("model", "grad_accumulator")
self._jitted_fwd_bwd_step_fn = nnx.jit(
fwd_bwd_step, donate_argnames=donate_argnames,
fwd_bwd_step,
donate_argnames=donate_argnames,
)
self._jitted_update_step_fn = nnx.jit(
update_step, donate_argnames=("optimizer", "grad_accumulator")
Expand Down Expand Up @@ -1068,7 +1070,11 @@ def _record_update(self, grad_norm: ArrayLike) -> int:
@override
def fwd_bwd(self, payload: datatypes.TrainerPayload | Any, **kwargs) -> None:
"""Executes forward and backward passes."""
fwd_bwd_step, _, _ = self.jit_fwd_bwd_update_and_eval_step()
cache_nnx_graph = kwargs.pop("cache_nnx_graph", True)
skip_jit = kwargs.pop("skip_jit", False)
fwd_bwd_step, _, _ = self.jit_fwd_bwd_update_and_eval_step(
skip_jit, cache_nnx_graph
)
self._record_fwd_bwd(
*fwd_bwd_step(
grad_accumulator=self.grad_accumulator,
Expand All @@ -1079,7 +1085,11 @@ def fwd_bwd(self, payload: datatypes.TrainerPayload | Any, **kwargs) -> None:
@override
def update(self, **kwargs) -> int:
"""Applies the accumulated gradients."""
_, update_step, _ = self.jit_fwd_bwd_update_and_eval_step()
cache_nnx_graph = kwargs.pop("cache_nnx_graph", True)
skip_jit = kwargs.pop("skip_jit", False)
_, update_step, _ = self.jit_fwd_bwd_update_and_eval_step(
skip_jit, cache_nnx_graph
)
return self._record_update(update_step())

def train_step(
Expand All @@ -1093,7 +1103,9 @@ def train_step(
available in the single-microstep regime; when accumulating there is work
between the two halves, so they must stay separate.
"""
self.jit_fwd_bwd_update_and_eval_step()
cache_nnx_graph = kwargs.pop("cache_nnx_graph", True)
skip_jit = kwargs.pop("skip_jit", False)
self.jit_fwd_bwd_update_and_eval_step(skip_jit, cache_nnx_graph)
if self._jitted_train_step_fn is None:
raise ValueError(
"train_step() requires exactly one micro-batch per update. Use"
Expand Down Expand Up @@ -1121,7 +1133,11 @@ def eval_step(
a sequence of eval_step calls with eval_context() so that the metrics
mode is set to EVAL and buffered metrics are written on exit.
"""
_, _, eval_step_fn = self.jit_fwd_bwd_update_and_eval_step()
skip_jit = kwargs.pop("skip_jit", False)
cache_nnx_graph = kwargs.pop("cache_nnx_graph", True)
_, _, eval_step_fn = self.jit_fwd_bwd_update_and_eval_step(
skip_jit=skip_jit, cache_nnx_graph=cache_nnx_graph,
)
loss, aux = eval_step_fn(self._prepare_payload(payload))
loss = jax.lax.stop_gradient(loss)
self._buffered_eval_metrics = self._buffer_metrics(
Expand Down Expand Up @@ -1302,9 +1318,7 @@ def prepare_weight_sync(self, sync_request: Any = None, **kwargs) -> Any:
worker = _default_weight_sync_worker()
self._weight_sync_worker = worker

backend = (
"vllm_jax" if "vllm" in self._sampler_type else self._sampler_type
)
backend = "vllm_jax" if "vllm" in self._sampler_type else self._sampler_type
mapping_config = getattr(self.config, "mapping_config", None)
if (
mapping_config is None
Expand All @@ -1313,6 +1327,7 @@ def prepare_weight_sync(self, sync_request: Any = None, **kwargs) -> Any:
):
try:
from tunix.generate import mappings as mappings_lib # pylint: disable=g-import-not-at-top

mapping_config = mappings_lib.MappingConfig.build(
model=self.model, backend=backend
)
Expand All @@ -1325,6 +1340,7 @@ def prepare_weight_sync(self, sync_request: Any = None, **kwargs) -> Any:
and mapping_config.to_hf_mappings
):
from tunix.generate import utils as gen_utils # pylint: disable=g-import-not-at-top

converted_state = gen_utils.transfer_state_with_mappings(
src_state=nnx.state(self.model),
dst_state=self._target_state,
Expand Down Expand Up @@ -1501,15 +1517,23 @@ def train(
tags=tags,
) as span_v2:
if self._jitted_train_step_fn is not None and is_update_step_val:
self.train_step(train_example)
self.train_step(
train_example,
skip_jit=skip_jit,
cache_nnx_graph=cache_nnx_graph,
)
computation_to_track = self._last_update_grad_norm
else:
self.fwd_bwd(train_example)
self.fwd_bwd(
train_example,
skip_jit=skip_jit,
cache_nnx_graph=cache_nnx_graph,
)
assert self._buffered_train_metrics is not None
train_loss = self._buffered_train_metrics.losses[-1]
computation_to_track = train_loss
if is_update_step_val:
self.update()
self.update(skip_jit=skip_jit, cache_nnx_graph=cache_nnx_graph)
computation_to_track = getattr(
self, "_last_update_grad_norm", train_loss
)
Expand Down
Loading