[SFT] Cut optimizer-state HBM by skipping depth-1 accumulator and keeping moment dtype.

tests/sft/peft_trainer_test.py

@@ -134,6 +134,38 @@ def _train_step(
       trainer.train(self.train_ds, self.eval_ds)
     self.assertEqual(global_counter, 1)
 
+  def test_compile_once_on_cond_path(self):
+    """Depth>1 (accumulator + nnx.cond path) also traces exactly once.
+
+    Guards the moment-dtype change (bf16 by default on the cond path) against
+    re-tracing: the cond path must still compile just once.
+    """
+    class CountCompiledTimesTrainer(peft_trainer.PeftTrainer):
+
+      def _train_step(
+          self, model, optimizer, grad_accumulator, inputs, is_update_step
+      ):
+        global global_counter
+        global_counter += 1
+        return super()._train_step(
+            model, optimizer, grad_accumulator, inputs, is_update_step
+        )
+
+    config = peft_trainer.TrainingConfig(
+        eval_every_n_steps=2, max_steps=100, gradient_accumulation_steps=2
+    )
+    rngs = nnx.Rngs(0)
+    model = tc.get_lora_model(
+        tc.ToyTransformer(config=tc.ModelConfig(), rngs=rngs), mesh=self.mesh
+    )
+    trainer = CountCompiledTimesTrainer(model, optax.sgd(1e-3), config)
+    trainer = trainer.with_gen_model_input_fn(dummy_gen_model_input_fn)
+    global global_counter
+    global_counter = 0
+    with self.mesh:
+      trainer.train(self.train_ds, self.eval_ds)
+    self.assertEqual(global_counter, 1)
+
   @parameterized.named_parameters(
       ('cache_nnx_graph', True),
       ('no_cache_nnx_graph', False),
@@ -1331,13 +1363,12 @@ def step(model, optimizer, acc, is_update_step):
     for d in float_dtypes:
       self.assertEqual(d, jnp.bfloat16)
 
-  def test_peft_trainer_promotes_bf16_opt_state_floats_to_float32(self):
-    """`PeftTrainer.__init__` casts float opt_state leaves to float32.
+  def test_peft_trainer_keeps_bf16_moments_on_cond_path(self):
+    """Default None keeps bf16 moments on the cond path, even with an fp32
 
-    `optax.adam` / `optax.adamw` promote their floating-point moments
-    (`mu`, `nu`) to float32 inside `update` whenever the learning rate is
-    a float32 tracer (as produced by `optax.inject_hyperparams`). This test
-    verifies that the trainer casts these to float32 in-place during init.
+    inject_hyperparams learning rate. Depth>1 (and packing) take the `nnx.cond`
+    path; bf16 moments trace fine there, so nothing is promoted to float32. Set
+    `optimizer_state_dtype=jnp.float32` to force fp32 moments.
     """
     rngs = nnx.Rngs(0)
     model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=rngs)
@@ -1352,7 +1383,9 @@ def test_peft_trainer_promotes_bf16_opt_state_floats_to_float32(self):
     tx = optax.inject_hyperparams(optax.adamw, hyperparam_dtype=jnp.float32)(
         learning_rate=1e-3
     )
-    config = peft_trainer.TrainingConfig(eval_every_n_steps=100, max_steps=1)
+    config = peft_trainer.TrainingConfig(
+        eval_every_n_steps=100, max_steps=1, gradient_accumulation_steps=2
+    )  # depth>1 -> cond path; moments stay bf16 (param dtype)
     trainer = peft_trainer.PeftTrainer(model, tx, config)
 
     opt_state_dtypes = jax.tree_util.tree_leaves(
@@ -1362,12 +1395,13 @@ def test_peft_trainer_promotes_bf16_opt_state_floats_to_float32(self):
             is_leaf=lambda x: isinstance(x, nnx.Variable),
         )
     )
-    float_dtypes = [
-        d for d in opt_state_dtypes if jnp.issubdtype(d, jnp.floating)
-    ]
-    self.assertNotEmpty(float_dtypes)
-    for d in float_dtypes:
-      self.assertEqual(d, jnp.float32)
+    float_dtypes = {
+        str(d) for d in opt_state_dtypes if jnp.issubdtype(d, jnp.floating)
+    }
+    # Moments (mu/nu, the bulk of opt-state) stay bf16 — not promoted to fp32.
+    # The only fp32 float leaf is the injected learning rate (hyperparam_dtype).
+    self.assertIn('bfloat16', float_dtypes)
+    self.assertEqual(float_dtypes - {'bfloat16'}, {'float32'})
 
 
 class Depth1FastPathTest(parameterized.TestCase):
@@ -1550,5 +1584,89 @@ def test_depth2_cadence(self):
     np.testing.assert_array_equal(trainer.grad_accumulator.denom[...], 0.0)
 
 
+class OptimizerMemoryTest(parameterized.TestCase):
+  """Optimizer-state HBM knobs: moment dtype + depth-1 accumulator skip."""
+
+  def _bf16_model(self):
+    model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
+    state = nnx.state(model, nnx.Param)
+    state = jax.tree.map(
+        lambda v: v.astype(jnp.bfloat16)
+        if jnp.issubdtype(v.dtype, jnp.floating)
+        else v,
+        state,
+    )
+    nnx.update(model, state)
+    return model
+
+  def _moment_float_dtypes(self, trainer):
+    st = nnx.state(trainer.optimizer, nnx.optimizer.OptState)
+    return {
+        str(x.dtype)
+        for x in jax.tree_util.tree_leaves(st)
+        if hasattr(x, 'dtype') and jnp.issubdtype(x.dtype, jnp.floating)
+    }
+
+  @parameterized.named_parameters(
+      ('float32', jnp.float32, 'float32'),
+      ('bfloat16', jnp.bfloat16, 'bfloat16'),
+  )
+  def test_optimizer_state_dtype_casts_moments(self, dtype, expected):
+    """`optimizer_state_dtype` casts the Adam moment trees to that dtype."""
+    model = self._bf16_model()
+    config = peft_trainer.TrainingConfig(
+        eval_every_n_steps=100, max_steps=1, optimizer_state_dtype=dtype
+    )
+    trainer = peft_trainer.PeftTrainer(model, optax.adamw(1e-3), config)
+    self.assertEqual(self._moment_float_dtypes(trainer), {expected})
+
+  def test_optimizer_state_dtype_none_keeps_param_dtype_on_fast_path(self):
+    """Default None + depth-1 keeps moments at the param dtype (no forced fp32)."""
+    model = self._bf16_model()  # bf16 params
+    config = peft_trainer.TrainingConfig(eval_every_n_steps=100, max_steps=1)
+    trainer = peft_trainer.PeftTrainer(model, optax.adamw(1e-3), config)
+    # bf16 params -> bf16 moments (matching optax), NOT promoted to fp32.
+    self.assertEqual(self._moment_float_dtypes(trainer), {'bfloat16'})
+
+  def test_optimizer_state_dtype_none_keeps_param_dtype_on_cond_path(self):
+    """Default None + depth>1 keeps moments at the param dtype (no forced fp32)."""
+    model = self._bf16_model()  # bf16 params
+    config = peft_trainer.TrainingConfig(
+        eval_every_n_steps=100, max_steps=1, gradient_accumulation_steps=2
+    )
+    trainer = peft_trainer.PeftTrainer(model, optax.adamw(1e-3), config)
+    # bf16 params -> bf16 moments on the cond path too (set fp32 to override).
+    self.assertEqual(self._moment_float_dtypes(trainer), {'bfloat16'})
+
+  def test_accumulator_grads_skipped_at_depth1(self):
+    """At depth-1 (non-packing) the accumulator grad tree is not allocated."""
+    model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
+    config = peft_trainer.TrainingConfig(
+        eval_every_n_steps=100, max_steps=1
+    )  # gradient_accumulation_steps=None -> depth 1
+    trainer = peft_trainer.PeftTrainer(model, optax.sgd(1e-3), config)
+    self.assertEmpty(
+        jax.tree_util.tree_leaves(trainer.grad_accumulator.grads)
+    )
+
+  @parameterized.named_parameters(
+      ('depth2', 2, None),
+      ('packing', None, 16),
+  )
+  def test_accumulator_grads_allocated_when_used(self, steps, max_seq_token):
+    """Depth>1 or packing keeps the accumulator grad tree allocated."""
+    model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
+    config = peft_trainer.TrainingConfig(
+        eval_every_n_steps=100,
+        max_steps=1,
+        gradient_accumulation_steps=steps,
+        max_seq_token_per_tpu=max_seq_token,
+    )
+    trainer = peft_trainer.PeftTrainer(model, optax.sgd(1e-3), config)
+    self.assertNotEmpty(
+        jax.tree_util.tree_leaves(trainer.grad_accumulator.grads)
+    )
+
+
 if __name__ == '__main__':
   absltest.main()

tunix/sft/peft_trainer.py

@@ -29,6 +29,7 @@
 import jax.numpy as jnp
 import jax.sharding as shd
 from jax.typing import ArrayLike  # pylint: disable=g-importing-member
+from jax.typing import DTypeLike  # pylint: disable=g-importing-member
 import numpy as np
 import optax
 from tunix.perf import metrics as perf_metrics
@@ -89,6 +90,10 @@ class TrainingConfig:
   # Sequence packing configuration.
   max_seq_token_per_tpu: int | None = None
 
+  # Adam moment dtype. None (default) follows the param dtype (optax inits
+  # moments as zeros_like(params)); set e.g. jnp.float32 to force fp32.
+  optimizer_state_dtype: DTypeLike | None = None
+
   def get_with_default(self, key: str, default: Any) -> Any:
     val = getattr(self, key)
     if val is None:
@@ -189,9 +194,21 @@ class GradientAccumulator(nnx.Module):
     sizes.
   """
 
-  def __init__(self, model: nnx.Module, wrt: type[nnx.Variable]):
+  def __init__(
+      self,
+      model: nnx.Module,
+      wrt: type[nnx.Variable],
+      *,
+      allocate_grads: bool = True,
+  ):
     state = nnx.state(model, wrt)
-    self.grads = nnx.data(jax.tree_util.tree_map(jnp.zeros_like, state))
+    if allocate_grads:
+      self.grads = nnx.data(jax.tree_util.tree_map(jnp.zeros_like, state))
+    else:
+      # Fast path never reads the accumulator: skip the model-sized grad-tree
+      # allocation. Empty grads keep it a valid tiny jit arg (signature and
+      # compilation unchanged).
+      self.grads = nnx.data({})
     self.denom = nnx.Variable(jnp.zeros((), dtype=jnp.float32))
 
   def add(self, grads: Any, denom: jax.Array | None = None):
@@ -237,22 +254,20 @@ def _zero_in_place(v):
     self.denom.set_value(jnp.zeros_like(self.denom[...]))
 
 
-def _promote_opt_state_floats_to_float32(optimizer: nnx.Optimizer) -> None:
-  """Cast the optimizer state's floating-point leaves to float32 in-place.
-
-  Args:
-    optimizer: The nnx.Optimizer instance whose state will be modified.
-  """
+def _cast_opt_state_floats(
+    optimizer: nnx.Optimizer, dtype: jnp.dtype
+) -> None:
+  """Cast the optimizer state's floating-point leaves to `dtype` in-place."""
 
   def _cast(v):
     if isinstance(v, nnx.Variable):
       val = v.value
       if (
           hasattr(val, "dtype")
           and jnp.issubdtype(val.dtype, jnp.floating)
-          and val.dtype != jnp.float32
+          and val.dtype != dtype
       ):
-        v.value = val.astype(jnp.float32)
+        v.value = val.astype(dtype)
 
   opt_state = nnx.state(optimizer, nnx.optimizer.OptState)
   jax.tree_util.tree_map(
@@ -309,11 +324,20 @@ def __init__(
     self._lora_enabled = utils.is_lora_enabled(self.model)
     wrt_target = nnx.LoRAParam if self._lora_enabled else nnx.Param
     self.optimizer = nnx.Optimizer(self.model, optimizer, wrt=wrt_target)
-    # Promote floating-point leaves to float32 in-place to match the dtype of
-    # the optimizer update function branch (which is float32 due to
-    # `optax.inject_hyperparams`).
-    _promote_opt_state_floats_to_float32(self.optimizer)
-    self.grad_accumulator = GradientAccumulator(self.model, wrt_target)
+    # Adam moments follow the param dtype by default (optax inits them as
+    # zeros_like(params)). Set optimizer_state_dtype to override, e.g.
+    # jnp.float32.
+    if self.config.optimizer_state_dtype is not None:
+      _cast_opt_state_floats(self.optimizer, self.config.optimizer_state_dtype)
+    # Depth-1 non-packing fast path never reads the accumulator; skip its
+    # model-sized grad-tree allocation there.
+    _uses_cond_path = not (
+        self.config.get_with_default("gradient_accumulation_steps", 1) == 1
+        and self.config.max_seq_token_per_tpu is None
+    )
+    self.grad_accumulator = GradientAccumulator(
+        self.model, wrt_target, allocate_grads=_uses_cond_path
+    )
 
     self.loss_fn = _default_loss_fn
     self.eval_loss_fn = _default_loss_fn