megatron.checkpointing.fix_query_key_value_ordering#

megatron.checkpointing.fix_query_key_value_ordering(model, checkpoint_version)#

Fix up query/key/value matrix ordering if checkpoint version is smaller than 2.0