save and load quantized checkpoint #671
Conversation
MaxText/checkpointing.py
Outdated
| # ) | ||
|
|
||
| restored = ckptr.restore(p) | ||
| print(restored['params'].keys()) # printed ['params', 'aqt'] |
There was a problem hiding this comment.
@singh-mitali The original restore call can only load the 'params' key, but the saved checkpoint we currently have puts the weights in both 'params' and 'aqt'
This change can load the checkpoint with both 'params' and 'aqt'. But it has some sharding issue when actually running. I guess it's because we are missing the second argument items as in the original restore() call. Not sure how to split the abstract_unboxed_pre_state.params to match our new quantized checkpoint
There was a problem hiding this comment.
Another way we probably can save the quantized checkpoint as the original format so we don't need to change the loading code, but I found path keys of quantized checkpoint are different under 'aqt' compared to the original checkpoint under 'params'
There was a problem hiding this comment.
Yes they are different - they are generated by the aqt code and same paths will be expected by aqt code to read once we load the checkpoint. We would have to do the conversion on both ends in that case (also these paths can change due to aqt code changes or which layers we apply aqt to). There are the int8 values for weights but also an additional parameter scale per quantized vector. Better to see if we can just save the two parameters separately in checkpoint.
There was a problem hiding this comment.
So could create a state called QuantizedDecodeState which has both params and aqt as variables. Store that in checkpoint. Alternatively use multi-object checkpointing https://colab.sandbox.google.com/github/google/orbax/blob/main/checkpoint/orbax/checkpoint/orbax_checkpoint.ipynb
| state.params['params'] = params['params'] | ||
| state.params['aqt'] = params['aqt'] | ||
| if save_checkpoint(checkpoint_manager, step_number_to_save_new_ckpt, state): | ||
| print('save checkpoint successfully') |
There was a problem hiding this comment.
This can successfully save a quantized checkpoint
No description provided.