gm.ckpts.SkipLoRA

gm.ckpts.SkipLoRA#

class gemma.gm.ckpts.SkipLoRA(
wrapped: kauldron.checkpoints.partial_loader.InitTransform,
)[source]

Bases: kauldron.checkpoints.partial_loader.InitTransform

Wraps a partial loader to not restore the LoRA weights.

wrapped: kauldron.checkpoints.partial_loader.InitTransform
transform(
state: gemma.gm.ckpts._lora._StateT,
) → gemma.gm.ckpts._lora._StateT[source]

Transform the state by updating it with pre-trained values.

Notes:

  • transform functions can modify the state values but should NOT modify its structure, shape or dtypes.

  • transform should correctly propagate the sharding information from the given state.

Parameters:

state – The state object to transform

Returns:

The updated state