declare-lab/tango2
92
1# coding=utf-8
2# Copyright 2023 The HuggingFace Inc. team.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15""" PyTorch - Flax general utilities."""
16
17from pickle import UnpicklingError
18
19import jax
20import jax.numpy as jnp
21import numpy as np
22from flax.serialization import from_bytes
23from flax.traverse_util import flatten_dict
24
25from ..utils import logging
26
27
28logger = logging.get_logger(__name__)
29
30
31#####################
32# Flax => PyTorch #
33#####################
34
35
36# from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flax_pytorch_utils.py#L224-L352
37def load_flax_checkpoint_in_pytorch_model(pt_model, model_file):
38 try:
39 with open(model_file, "rb") as flax_state_f:
40 flax_state = from_bytes(None, flax_state_f.read())
41 except UnpicklingError as e:
42 try:
43 with open(model_file) as f:
44 if f.read().startswith("version"):
45 raise OSError(
46 "You seem to have cloned a repository without having git-lfs installed. Please"
47 " install git-lfs and run `git lfs install` followed by `git lfs pull` in the"
48 " folder you cloned."
49 )
50 else:
51 raise ValueError from e
52 except (UnicodeDecodeError, ValueError):
53 raise EnvironmentError(f"Unable to convert {model_file} to Flax deserializable object. ")
54
55 return load_flax_weights_in_pytorch_model(pt_model, flax_state)
56
57
58def load_flax_weights_in_pytorch_model(pt_model, flax_state):
59 """Load flax checkpoints in a PyTorch model"""
60
61 try:
62 import torch # noqa: F401
63 except ImportError:
64 logger.error(
65 "Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"
66 " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"
67 " instructions."
68 )
69 raise
70
71 # check if we have bf16 weights
72 is_type_bf16 = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype == jnp.bfloat16, flax_state)).values()
73 if any(is_type_bf16):
74 # convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16
75
76 # and bf16 is not fully supported in PT yet.
77 logger.warning(
78 "Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` "
79 "before loading those in PyTorch model."
80 )
81 flax_state = jax.tree_util.tree_map(
82 lambda params: params.astype(np.float32) if params.dtype == jnp.bfloat16 else params, flax_state
83 )
84
85 pt_model.base_model_prefix = ""
86
87 flax_state_dict = flatten_dict(flax_state, sep=".")
88 pt_model_dict = pt_model.state_dict()
89
90 # keep track of unexpected & missing keys
91 unexpected_keys = []
92 missing_keys = set(pt_model_dict.keys())
93
94 for flax_key_tuple, flax_tensor in flax_state_dict.items():
95 flax_key_tuple_array = flax_key_tuple.split(".")
96
97 if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
98 flax_key_tuple_array = flax_key_tuple_array[:-1] + ["weight"]
99 flax_tensor = jnp.transpose(flax_tensor, (3, 2, 0, 1))
100 elif flax_key_tuple_array[-1] == "kernel":
101 flax_key_tuple_array = flax_key_tuple_array[:-1] + ["weight"]
102 flax_tensor = flax_tensor.T
103 elif flax_key_tuple_array[-1] == "scale":
104 flax_key_tuple_array = flax_key_tuple_array[:-1] + ["weight"]
105
106 if "time_embedding" not in flax_key_tuple_array:
107 for i, flax_key_tuple_string in enumerate(flax_key_tuple_array):
108 flax_key_tuple_array[i] = (
109 flax_key_tuple_string.replace("_0", ".0")
110 .replace("_1", ".1")
111 .replace("_2", ".2")
112 .replace("_3", ".3")
113 )
114
115 flax_key = ".".join(flax_key_tuple_array)
116
117 if flax_key in pt_model_dict:
118 if flax_tensor.shape != pt_model_dict[flax_key].shape:
119 raise ValueError(
120 f"Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected "
121 f"to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."
122 )
123 else:
124 # add weight to pytorch dict
125 flax_tensor = np.asarray(flax_tensor) if not isinstance(flax_tensor, np.ndarray) else flax_tensor
126 pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
127 # remove from missing keys
128 missing_keys.remove(flax_key)
129 else:
130 # weight is not expected by PyTorch model
131 unexpected_keys.append(flax_key)
132
133 pt_model.load_state_dict(pt_model_dict)
134
135 # re-transform missing_keys to list
136 missing_keys = list(missing_keys)
137
138 if len(unexpected_keys) > 0:
139 logger.warning(
140 "Some weights of the Flax model were not used when initializing the PyTorch model"
141 f" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"
142 f" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"
143 " (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"
144 f" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"
145 " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"
146 " FlaxBertForSequenceClassification model)."
147 )
148 if len(missing_keys) > 0:
149 logger.warning(
150 f"Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"
151 f" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"
152 " use it for predictions and inference."
153 )
154
155 return pt_model
156