Text Generation
Transformers
PyTorch
Safetensors
Finnish
llama
finnish
text-generation-inference
llama-7b-finnish / EasyLM /scripts /diff_checkpoint.py
aapot
Add files
a971b09
raw
history blame
1.85 kB
# This script converts model checkpoint trained by EsayLM to a standard
# mspack checkpoint that can be loaded by huggingface transformers or
# flax.serialization.msgpack_restore. Such conversion allows models to be
# used by other frameworks that integrate with huggingface transformers.
import pprint
from functools import partial
import os
import numpy as np
import jax
import jax.numpy as jnp
import flax.serialization
import mlxu
from EasyLM.checkpoint import StreamingCheckpointer
from EasyLM.jax_utils import float_to_dtype
FLAGS, FLAGS_DEF = mlxu.define_flags_with_default(
recover_diff=False,
load_base_checkpoint='',
load_target_checkpoint='',
output_file='',
streaming=True,
float_dtype='bf16',
)
def main(argv):
assert FLAGS.load_base_checkpoint != '' and FLAGS.load_target_checkpoint != ''
assert FLAGS.output_file != ''
base_params = StreamingCheckpointer.load_trainstate_checkpoint(
FLAGS.load_base_checkpoint, disallow_trainstate=True
)[1]['params']
target_params = StreamingCheckpointer.load_trainstate_checkpoint(
FLAGS.load_target_checkpoint, disallow_trainstate=True
)[1]['params']
if FLAGS.recover_diff:
params = jax.tree_util.tree_map(
lambda b, t: b + t, base_params, target_params
)
else:
params = jax.tree_util.tree_map(
lambda b, t: t - b, base_params, target_params
)
if FLAGS.streaming:
StreamingCheckpointer.save_train_state_to_file(
params, FLAGS.output_file, float_dtype=FLAGS.float_dtype
)
else:
params = float_to_dtype(params, FLAGS.float_dtype)
with mlxu.open_file(FLAGS.output, 'wb') as fout:
fout.write(flax.serialization.msgpack_serialize(params, in_place=True))
if __name__ == "__main__":
mlxu.run(main)