Spaces:
Build error
Build error
File size: 3,034 Bytes
b100e1c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
# Copyright 2022 The T5X Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for training_utils."""
import functools
import os
# Emulate 2 devices on CPU. Import before JAX.
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=2'
from absl.testing import absltest # pylint: disable=g-import-not-at-top
from flax import core as flax_core
import jax
from jax import numpy as jnp
import numpy as np
from t5x.contrib.moe import training_utils
class MatchFnTest(absltest.TestCase):
def test_regex_prefix(self):
match_fn = training_utils.match_fn(r'.*test.*')
self.assertTrue(match_fn('/test/something'))
self.assertTrue(match_fn('to/test/or/not/'))
self.assertFalse(match_fn('no/match'))
def test_empty_prefix(self):
match_fn = training_utils.match_fn(None)
self.assertFalse(match_fn('/test/something'))
self.assertFalse(match_fn('to/test/or/not/'))
class ScaleShardedGradsTest(absltest.TestCase):
def test_scale_sharded_grads(self):
grads = flax_core.freeze({
'encoder': {
'expert_layer': jnp.ones((2, 3)),
'regular_layer': jnp.ones((1, 2))
}
})
sharded_match_fn = training_utils.match_fn(r'.*expert.*')
scaled_grads = training_utils.scale_sharded_grads(
grads, sharded_match_fn, scale_factor=100.)
expected_grads = flax_core.freeze({
'encoder': {
'expert_layer': 100. * jnp.ones((2, 3)),
'regular_layer': jnp.ones((1, 2))
}
})
jax.tree_map(
functools.partial(np.testing.assert_allclose, rtol=3e-7), scaled_grads,
expected_grads)
class TreeTest(absltest.TestCase):
def test_tree_flatten_with_names(self):
tree = {'ff_0': {'kernel': 0, 'bias': 1}, 'ff_1': {'kernel': 2, 'bias': 3}}
names_and_values, _ = training_utils._tree_flatten_with_names(tree)
expected_names_and_values = [('ff_0/bias', 1), ('ff_0/kernel', 0),
('ff_1/bias', 3), ('ff_1/kernel', 2)]
self.assertEqual(names_and_values, expected_names_and_values)
# Check that values match regular JAX tree_flatten.
self.assertEqual([x for _, x in names_and_values],
jax.tree_flatten(tree)[0])
def test_tree_map_with_names(self):
tree = {'a': 1, 'b': 2}
mapped_tree = training_utils.tree_map_with_names(
f=lambda x: -x, param_tree=tree, match_name_fn=lambda name: name == 'b')
self.assertEqual(mapped_tree, {'a': 1, 'b': -2})
if __name__ == '__main__':
absltest.main()
|