import re from flax.core.frozen_dict import freeze from flax.traverse_util import flatten_dict, unflatten_dict from jax.experimental import PartitionSpec as P # utils adapted from https://github.com/google-research/google-research/blob/master/flax_models/t5x/partitions.py # Sentinels _unmatched = object() # For specifying empty leaf dict `{}` empty_dict = object() def _match(qs, ks): """Return True if regexes in qs match any window of strings in tuple ks.""" # compile regexes and force complete match qts = tuple(map(lambda x: re.compile(x + "$"), qs)) for i in range(len(ks) - len(qs) + 1): matches = [x.match(y) for x, y in zip(qts, ks[i:])] if matches and all(matches): return True return False def _replacement_rules(rules): def replace(key, val): for rule, replacement in rules: if _match(rule, key): return replacement return val return replace def _get_partition_rules(): return [ # embeddings (("embed_positions", "embedding"), P("mp", None)), (("embed_tokens", "embedding"), P("mp", None)), (("rel_bias", "embedding"), P(None, "mp")), # attention (("(q_proj|k_proj|v_proj)", "kernel"), P(None, "mp")), (("out_proj", "kernel"), P("mp", None)), # FFN (("Dense_0", "kernel"), P(None, "mp")), (("GLU.*", "Dense_1", "kernel"), P(None, "mp")), (("GLU.*", "Dense_2", "kernel"), P("mp", None)), (("FFN.*", "Dense_1", "kernel"), P("mp", None)), # layer norms (("(bias|scale)",), None), (("lm_head", "kernel"), P(None, "mp")), # head scale and tau (("(head_scale|tau)",), None), ] def set_partitions(in_dict): rules = _get_partition_rules() replace = _replacement_rules(rules) initd = {k: _unmatched for k in flatten_dict(in_dict)} result = {k: replace(k, v) for k, v in initd.items()} for k, v in result.items(): if v == _unmatched: print(f"Unmatched -> {k}") assert _unmatched not in result.values(), "Incomplete partition spec." return freeze(unflatten_dict(result))