How does Tunix freeze the base model during a LoRA fine-tune?

By never handing it to the optimizer. The trainer walks the model graph once looking for a single nnx.LoRAParam leaf; if it finds one, it sets wrt_target to nnx.LoRAParam and builds nnx.Optimizer(self.model, optimizer, wrt=wrt_target). Optimizer state is created only for leaves that match that filter, so a base weight has no moment vectors to update. The same type is passed as nnx.DiffState(0, nnx.LoRAParam) to nnx.value_and_grad, so the backbone is not differentiated either. There is no requires_grad flag and no mask: the freeze is the absence of state, not a switch on the parameter.

Is the PeftTrainer only for LoRA?

No. The class docstring says it is a PEFT trainer for LoRA, but the LoRA-specific behaviour is one boolean. utils.is_lora_enabled returns False on a model with no adapter leaves, wrt_target becomes nnx.Param, and the same train step does a full fine-tune of every parameter. Checkpoint save and restore flip on the same boolean. Everything else in the file, the gradient accumulator, the metrics buffer, the eval loop, the jit and donation setup, runs identically in both modes.

How does gradient accumulation work in the Tunix PEFT trainer?

Through a GradientAccumulator module, not optax.MultiSteps. It allocates a float32 zeros buffer shaped like nnx.state(model, wrt) plus one scalar denominator. Each micro-step adds its gradients and adds to the denominator, and get() divides the sum by the sum of denominators, then casts back to the parameter dtype. Passing the real denominator matters when a micro-batch has an unequal number of valid target tokens: the result is the global weighted mean rather than a mean of means. When gradient_accumulation_steps is 1 and sequence packing is off, the train step skips the accumulator entirely and updates from the gradients directly, and the accumulator is constructed with an empty buffer to save memory.

What does a LoRA checkpoint contain in Tunix?

Only the adapter parameters and whatever optimizer state exists. The trainer passes save_only_lora_params=self._lora_enabled to the checkpoint manager, which narrows the model state to nnx.state(model, nnx.LoRAParam). The optimizer state is not narrowed by that flag, but it does not need to be: the optimizer was built over the same filter, so it holds nothing else. On resume the trainer calls maybe_restore with restore_only_lora_params set the same way, and the restore path turns on Orbax partial loading so the missing base weights are not treated as an error.

Why does the trainer log metrics one step late?

To overlap the write with the next step. _write_train_metrics keeps two buffers: it holds the step that just finished and writes the previous one. Reading a metric forces the JAX arrays for that step to be materialised on the host, which would otherwise block dispatch of the next step, so the first call writes nothing and every later call writes the step before. It also increments the buffered step by one before writing, because the train-step counter is not incremented until the update actually lands.

The Tunix PEFT Trainer, Line by Line

A LoRA fine-tune freezes the base model by leaving it out. This walk reads tunix/sft/peft_trainer.py at one pinned commit: the type tag that decides which parameters an optimizer is even built over, the two update paths inside one train step, the gradient accumulator that carries a denominator, and the checkpoint that writes only the adapter.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of The Tunix PEFT Trainer, Line by Line — built to be understood, not skimmed.

Questions

How does Tunix freeze the base model during a LoRA fine-tune?: By never handing it to the optimizer. The trainer walks the model graph once looking for a single nnx.LoRAParam leaf; if it finds one, it sets wrt_target to nnx.LoRAParam and builds nnx.Optimizer(self.model, optimizer, wrt=wrt_target). Optimizer state is created only for leaves that match that filter, so a base weight has no moment vectors to update. The same type is passed as nnx.DiffState(0, nnx.LoRAParam) to nnx.value_and_grad, so the backbone is not differentiated either. There is no requires_grad flag and no mask: the freeze is the absence of state, not a switch on the parameter.
Is the PeftTrainer only for LoRA?: No. The class docstring says it is a PEFT trainer for LoRA, but the LoRA-specific behaviour is one boolean. utils.is_lora_enabled returns False on a model with no adapter leaves, wrt_target becomes nnx.Param, and the same train step does a full fine-tune of every parameter. Checkpoint save and restore flip on the same boolean. Everything else in the file, the gradient accumulator, the metrics buffer, the eval loop, the jit and donation setup, runs identically in both modes.
How does gradient accumulation work in the Tunix PEFT trainer?: Through a GradientAccumulator module, not optax.MultiSteps. It allocates a float32 zeros buffer shaped like nnx.state(model, wrt) plus one scalar denominator. Each micro-step adds its gradients and adds to the denominator, and get() divides the sum by the sum of denominators, then casts back to the parameter dtype. Passing the real denominator matters when a micro-batch has an unequal number of valid target tokens: the result is the global weighted mean rather than a mean of means. When gradient_accumulation_steps is 1 and sequence packing is off, the train step skips the accumulator entirely and updates from the gradients directly, and the accumulator is constructed with an empty buffer to save memory.
What does a LoRA checkpoint contain in Tunix?: Only the adapter parameters and whatever optimizer state exists. The trainer passes save_only_lora_params=self._lora_enabled to the checkpoint manager, which narrows the model state to nnx.state(model, nnx.LoRAParam). The optimizer state is not narrowed by that flag, but it does not need to be: the optimizer was built over the same filter, so it holds nothing else. On resume the trainer calls maybe_restore with restore_only_lora_params set the same way, and the restore path turns on Orbax partial loading so the missing base weights are not treated as an error.
Why does the trainer log metrics one step late?: To overlap the write with the next step. _write_train_metrics keeps two buffers: it holds the step that just finished and writes the previous one. Reading a metric forces the JAX arrays for that step to be materialised on the host, which would otherwise block dispatch of the next step, so the first call writes nothing and every later call writes the step before. It also increments the buffered step by one before writing, because the train-step counter is not incremented until the update actually lands.

The Tunix PEFT Trainer, Line by Line

Questions

Related explainers