def test_extract_prefixed_keys_from_state_specs(self):
w_sepc = base_layer.var_partition_specs(
{'w': py_utils.weight_params(shape=(4, 8))},
device_mesh=np.arange(1).reshape([1, 1]),
device_axis_names=['a', 'b'])
train_state_partition_specs = train_states.TrainState(
step=pjit.PartitionSpec(), mdl_vars=w_sepc, opt_states={})
nested_names = py_utils.extract_prefixed_keys_from_nested_map(
train_state_partition_specs)
flattened_names, _ = jax.tree_util.tree_flatten(nested_names)
self.assertListEqual(['step', 'mdl_vars/w'], flattened_names)
def test_multioptimizer_learner_adam_adagrad(self):
learner_p = learners.MultiOptimizerLearner.Params()
learner_p.name = 'learner'
learner_p.loss_name = 'loss'
learner_p.optimizer = optimizers.Adam.ParamsA()
learner_p.optimizer.learning_rate = 1.
learner_p.optimizer.lr_schedule = schedules.Constant.Params()
aux_p1 = optimizers.Adam.ParamsA()
aux_p1.lr_schedule = schedules.Constant.Params()
aux_p1.learning_rate = 2.0
aux_p2 = optimizers.Adagrad.Params()
aux_p2.lr_schedule = schedules.Constant.Params()
aux_p2.learning_rate = 3.0
learner_p.auxiliary_optimizers = [aux_p1, aux_p2]
learner_p.auxiliary_regex = ['ngram', 'transformer']
learner_instance = learner_p.Instantiate()
grads = NestedMap()
grads.lm = NestedMap()
grads.lm.ngrammer = NestedMap()
old_vars = grads.DeepCopy()
emb_var1 = jnp.asarray(
np.random.normal(1.0, 0.5, [4, 8]).astype('float32'))
emb_var2 = jnp.asarray(
np.random.normal(1.0, 0.5, [4, 8]).astype('float32'))
grad1 = jnp.asarray(
np.random.normal(1.0, 0.5, [4, 8]).astype('float32'))
grad2 = jnp.asarray(
np.random.normal(1.0, 0.5, [4, 8]).astype('float32'))
grads.lm.ngrammer.ngram_table = [
NestedMap(emb_var=grad1),
NestedMap(emb_var=grad2)
]
old_vars.lm.ngrammer.ngram_table = [
NestedMap(emb_var=emb_var1),
NestedMap(emb_var=emb_var2)
]
# Create some other keys.
grads.lm.transformer = NestedMap(w=jnp.asarray(
np.random.normal(1.4, 2.0, [4, 4]).astype('float32')))
old_vars.lm.transformer = NestedMap(w=jnp.asarray(
np.random.normal(1.2, 4.0, [4, 4]).astype('float32')))
grads.lm.ffn = NestedMap(k=jnp.asarray(
np.random.normal(1.6, 2.0, [4, 4]).astype('float32')))
old_vars.lm.ffn = NestedMap(k=jnp.asarray(
np.random.normal(1.3, 2.0, [4, 4]).astype('float32')))
var_weight_params = jax.tree_map(
lambda v: py_utils.weight_params(v.shape), old_vars)
grad_tx = learner_instance.get_grad_tx(var_weight_params)
opt_states = grad_tx.init(old_vars)
logging.info('opt_states: %s', opt_states)
def create_layer_variables(self):
super().create_layer_variables()
p = self.params
assert p.latent_dim % p.num_groups == 0
weight_params = py_utils.weight_params(
shape=[
p.num_latent_classes, p.num_groups, p.latent_dim // p.num_groups
],
init=p.params_init,
dtype=jnp.float32)
# [C, D]
self.create_variable('w', weight_params)
def test_multioptimizer_learner(self, lr_multiplier1, lr_multiplier2,
use_vq_ngrams):
learner_p = learners.MultiOptimizerLearner.Params()
learner_p.name = 'learner'
learner_p.loss_name = 'loss'
learner_p.optimizer = optimizers.Sgd.Params()
learner_p.optimizer.learning_rate = 1.
learner_p.optimizer.lr_schedule = schedules.Constant.Params()
aux_p1 = optimizers.Sgd.Params()
aux_p1.lr_schedule = schedules.Constant.Params()
aux_p1.learning_rate = lr_multiplier1
aux_p2 = optimizers.Sgd.Params()
aux_p2.lr_schedule = schedules.Constant.Params()
aux_p2.learning_rate = lr_multiplier2
learner_p.auxiliary_optimizers = [aux_p1, aux_p2]
learner_p.auxiliary_regex = ['ngram', 'transformer']
learner_instance = learner_p.Instantiate()
grads = NestedMap()
grads.lm = NestedMap()
grads.lm.ngrammer = NestedMap()
old_vars = grads.DeepCopy()
emb_var1 = jnp.asarray(
np.random.normal(1.0, 0.5, [4, 8]).astype('float32'))
emb_var2 = jnp.asarray(
np.random.normal(1.0, 0.5, [4, 8]).astype('float32'))
grad1 = jnp.asarray(
np.random.normal(1.0, 0.5, [4, 8]).astype('float32'))
grad2 = jnp.asarray(
np.random.normal(1.0, 0.5, [4, 8]).astype('float32'))
if use_vq_ngrams:
grads.lm.ngrammer.ngram_layer = NestedMap()
grads.lm.ngrammer.ngram_layer.ngram_table = [
NestedMap(emb_var=grad1),
NestedMap(emb_var=grad2)
]
old_vars.lm.ngrammer.ngram_layer = NestedMap()
old_vars.lm.ngrammer.ngram_layer.ngram_table = [
NestedMap(emb_var=emb_var1),
NestedMap(emb_var=emb_var2)
]
else:
grads.lm.ngrammer.ngram_table = [
NestedMap(emb_var=grad1),
NestedMap(emb_var=grad2)
]
old_vars.lm.ngrammer.ngram_table = [
NestedMap(emb_var=emb_var1),
NestedMap(emb_var=emb_var2)
]
# Create some other keys.
grads.lm.transformer = NestedMap(w=jnp.asarray(
np.random.normal(1.4, 2.0, [4, 4]).astype('float32')))
old_vars.lm.transformer = NestedMap(w=jnp.asarray(
np.random.normal(1.2, 4.0, [4, 4]).astype('float32')))
grads.lm.ffn = NestedMap(k=jnp.asarray(
np.random.normal(1.6, 2.0, [4, 4]).astype('float32')))
old_vars.lm.ffn = NestedMap(k=jnp.asarray(
np.random.normal(1.3, 2.0, [4, 4]).astype('float32')))
var_weight_params = jax.tree_map(
lambda v: py_utils.weight_params(v.shape), old_vars)
grad_tx = learner_instance.get_grad_tx(var_weight_params)
opt_states = grad_tx.init(old_vars)
with base_layer.JaxContext.new_context():
transformed_grads, _ = learner_instance.update_states(
grads, opt_states, old_vars, var_weight_params)
expected_grad1 = -lr_multiplier1 * grad1
expected_grad2 = -lr_multiplier1 * grad2
if use_vq_ngrams:
new_grad1 = (transformed_grads.lm.ngrammer.ngram_layer.
ngram_table[0].emb_var)
new_grad2 = (transformed_grads.lm.ngrammer.ngram_layer.
ngram_table[1].emb_var)
else:
new_grad1 = transformed_grads.lm.ngrammer.ngram_table[0].emb_var
new_grad2 = transformed_grads.lm.ngrammer.ngram_table[1].emb_var
self.assertAllClose(new_grad1, expected_grad1)
self.assertAllClose(new_grad2, expected_grad2)
expected_grad_transformer = -lr_multiplier2 * grads.lm.transformer.w
new_grad_transformer = transformed_grads.lm.transformer.w
expected_grad_ffn = -grads.lm.ffn.k
new_grad_ffn = transformed_grads.lm.ffn.k
self.assertAllClose(new_grad_transformer, expected_grad_transformer)
self.assertAllClose(new_grad_ffn, expected_grad_ffn)
def test_vectorized_prefix(self):
def _opt_init(params):
# Reduction over each variable. Behavior will depend on vectorization.
return jax.tree_map(jnp.sum, params)
def _opt_update(updates, state, params):
del params
return jax.tree_map(lambda u, s: u + s, updates, state), state
def _init_partition_spec(var_params):
def _init_one(p):
assert not p.repeat_prefix
return p
return jax.tree_map(_init_one, var_params)
class TestOptimizer(optimizers.BaseOptimizer):
def _get_raw_grad_transformation(self, lr):
return optimizers.ShardedGradientTransformation(
init=_opt_init,
update=_opt_update,
init_partition_spec=_init_partition_spec)
learner_p = learners.Learner.Params().Set(
name='learner', loss_name='loss', grad_norm_individual_vars=True)
learner_p.optimizer = TestOptimizer.Params().Set(
learning_rate=1., lr_schedule=schedules.Constant.Params())
learner_instance = learner_p.Instantiate()
grads = NestedMap(a=jnp.array([1, 2], dtype=jnp.float32),
b=jnp.array([1, 2], dtype=jnp.float32),
c=jnp.array([[1, 2], [3, 4]], dtype=jnp.float32))
variables = grads.copy()
a_var_param = py_utils.weight_params(())
a_var_param.repeat_prefix = [2]
a_var_param.repeat_prefix_split_dims_mapping = [-1]
b_var_param = py_utils.weight_params((2, ))
c_var_param = py_utils.weight_params(())
c_var_param.repeat_prefix = [2, 2]
c_var_param.repeat_prefix_split_dims_mapping = [('data', 'mdl'), None]
var_params = NestedMap(a=a_var_param, b=b_var_param, c=c_var_param)
grad_tx = learner_instance.get_grad_tx(var_weight_params=var_params)
partition_spec = grad_tx.init_partition_spec(var_params)
# Optimizers are chained as l1 - l2 - optimizer update - weight_decay.
opt_idx = 2
self.assertEqual(partition_spec['p#2#i-1'][opt_idx].a.shape, ())
self.assertEqual(partition_spec['p#2#i-1'][opt_idx].a.repeat_prefix,
[2])
self.assertEqual(
partition_spec['p#2#i-1']
[opt_idx].a.repeat_prefix_split_dims_mapping, [-1])
self.assertEqual(
partition_spec[opt_vec.NO_PREFIX_KEY][opt_idx].b.shape, (2, ))
self.assertEmpty(
partition_spec[opt_vec.NO_PREFIX_KEY][opt_idx].b.repeat_prefix
or [])
self.assertEqual(partition_spec['p#2.2#tsdata,smdl.'][opt_idx].c.shape,
())
self.assertEqual(
partition_spec['p#2.2#tsdata,smdl.'][opt_idx].c.repeat_prefix,
[2, 2])
self.assertEqual(
partition_spec['p#2.2#tsdata,smdl.']
[opt_idx].c.repeat_prefix_split_dims_mapping,
[('data', 'mdl'), None])
state = grad_tx.init(variables)
# Computed update is 0 + state, and state is sum of each variable.
update, state = grad_tx.update(jax.tree_map(jnp.zeros_like, variables),
state, variables)
# Variables a and c are scalars excluding the prefix, so the update must be
# equal to the initial variable values.
self.assertAllClose(update.a, variables.a)
self.assertAllClose(update.c, variables.c)
# b is not vectorized, so the update equals the sum reduction of the initial
# variable value.
self.assertAllClose(update.b,
jnp.zeros_like(variables.b) + jnp.sum(variables.b))
def test_multioptimizer_learner_sharding(self):
learner_p = learners.MultiOptimizerLearner.Params()
learner_p.name = 'learner'
learner_p.loss_name = 'loss'
learner_p.optimizer = optimizers.ShardedAdafactor.Params()
learner_p.optimizer.learning_rate = 1.
learner_p.optimizer.decay_method = 'pow'
learner_p.optimizer.lr_schedule = schedules.Constant.Params()
aux_p1 = optimizers.ShardedAdafactor.Params()
aux_p1.lr_schedule = schedules.Constant.Params()
aux_p1.learning_rate = 2.0
aux_p1.decay_method = 'pow'
aux_p2 = optimizers.ShardedAdafactor.Params()
aux_p2.lr_schedule = schedules.Constant.Params()
aux_p2.decay_method = 'adam'
aux_p2.learning_rate = 3.0
# Add auxiliary optimizers.
learner_p.auxiliary_optimizers = [aux_p1, aux_p2]
learner_p.auxiliary_regex = ['ngrammer', 'transformer']
learner_instance = learner_p.Instantiate()
# Add a single instance optimizer.
learner_p = learners.Learner.Params()
learner_p.name = 'learner'
learner_p.loss_name = 'loss'
learner_p.optimizer = optimizers.ShardedAdafactor.Params()
learner_p.optimizer.learning_rate = 1.
learner_p.optimizer.decay_method = 'pow'
learner_p.optimizer.lr_schedule = schedules.Constant.Params()
learner_instance_single = learner_p.Instantiate()
# Define device mesh.
mesh_shape = [1, 2, 1]
num_devices = np.prod(mesh_shape)
logging.info('num_local_devices: %s', num_devices)
device_mesh = np.arange(num_devices).reshape(mesh_shape)
grads = NestedMap()
grads.lm = NestedMap()
grads.lm.ngrammer = NestedMap()
old_vars = grads.DeepCopy()
emb_var1 = py_utils.weight_params(shape=[4, 8],
device_mesh=device_mesh,
tensor_split_dims_mapping=[-1, 1])
emb_var2 = py_utils.weight_params(shape=[4, 8],
device_mesh=device_mesh,
tensor_split_dims_mapping=[-1, 1])
grad1 = py_utils.weight_params(shape=[4, 8],
device_mesh=device_mesh,
tensor_split_dims_mapping=[-1, 1])
grad2 = py_utils.weight_params(shape=[4, 8],
device_mesh=device_mesh,
tensor_split_dims_mapping=[-1, 1])
grads.lm.ngrammer.ngram_table = [
NestedMap(emb_var=grad1),
NestedMap(emb_var=grad2)
]
old_vars.lm.ngrammer.ngram_table = [
NestedMap(emb_var=emb_var1),
NestedMap(emb_var=emb_var2)
]
# Create some other keys.
grads.lm.transformer = NestedMap(
w=py_utils.weight_params(shape=[4, 8],
device_mesh=device_mesh,
tensor_split_dims_mapping=[-1, 1]))
old_vars.lm.transformer = NestedMap(
w=py_utils.weight_params(shape=[4, 8],
device_mesh=device_mesh,
tensor_split_dims_mapping=[-1, 1]))
grads.lm.ffn = NestedMap(
k=py_utils.weight_params(shape=[4, 8],
device_mesh=device_mesh,
tensor_split_dims_mapping=[-1, 1]))
old_vars.lm.ffn = NestedMap(
k=py_utils.weight_params(shape=[4, 8],
device_mesh=device_mesh,
tensor_split_dims_mapping=[0, 1]))
grad_tx = learner_instance.get_grad_tx(var_weight_params=old_vars)
grad_tx_single = learner_instance_single.get_grad_tx(
var_weight_params=old_vars)
partition_spec = grad_tx.init_partition_spec(old_vars)
partition_spec_single = grad_tx_single.init_partition_spec(old_vars)
# Assert that the length of partition spec is the same as the total
# auxiliary optimizers plus 1 (for the primary optimizer).
self.assertLen(partition_spec,
len(learner_instance._auxiliary_optimizers) + 1)
# Optimizers are chained as l1 - l2 - optimizer update - weight_decay.
for k in partition_spec_single[2]._fields:
for p in partition_spec:
tf.nest.assert_same_structure(
getattr(p[2], k), getattr(partition_spec_single[2], k))