From 2b88d233fbecf1a065cff295c20ac04f370e84cc Mon Sep 17 00:00:00 2001 From: Lin Chai Date: Tue, 15 Sep 2026 11:51:25 -0700 Subject: [PATCH] [Tunix] Check on the `skip_jit` and `cache_nnx_graph` params from kwargs. `cache_nnx_graph` can save the cpu host time. PiperOrigin-RevId: 981947843 --- .../train/peft_trainer_v2_test.py | 20 +++--- tunix/experimental/train/peft_trainer_v2.py | 66 +++++++++++++------ 2 files changed, 54 insertions(+), 32 deletions(-) diff --git a/tests/experimental/train/peft_trainer_v2_test.py b/tests/experimental/train/peft_trainer_v2_test.py index 4270b2447d..2f89f232f0 100644 --- a/tests/experimental/train/peft_trainer_v2_test.py +++ b/tests/experimental/train/peft_trainer_v2_test.py @@ -789,8 +789,8 @@ 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, @@ -798,9 +798,7 @@ def _write_one_checkpoint(self, root, *, implicit_resume=True): '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 @@ -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) @@ -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): diff --git a/tunix/experimental/train/peft_trainer_v2.py b/tunix/experimental/train/peft_trainer_v2.py index 5b58984875..d78bc5b081 100644 --- a/tunix/experimental/train/peft_trainer_v2.py +++ b/tunix/experimental/train/peft_trainer_v2.py @@ -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`): @@ -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 @@ -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): @@ -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. @@ -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") @@ -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, @@ -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( @@ -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" @@ -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( @@ -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 @@ -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 ) @@ -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, @@ -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 )