Skip to content

Conversation

@Aman071106
Copy link

Title: docs(ckpts): Clarify 'quantize' behavior in load_params

Description

This PR addresses the TODO at line 277 regarding the quantize argument in _checkpoint.py.

Currently, the quantize=True argument in load_params does not perform actual weight quantization (e.g., to int8), but rather restructures the parameters for Quantization Aware Training (QAT). This can be misleading for users expecting immediate weight compression.

This change updates the docstring to explicitly state that quantize=True only enables the QAT structure, improving API clarity and developer experience.

Changes

  • gemma/gm/ckpts/_checkpoint.py: Updated load_params docstring to clarify that actual quantization happens during training, not loading.

Related Issue

Refers to: # TODO(epot): Better API. Currently this do not quantize the weights...

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant