google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
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

16
"""Tests for partitioning."""
17

18
import os
19

20
from absl.testing import absltest
21
import jax
22
from jax import core
23
import jax.numpy as jnp
24
from jax.sharding import NamedSharding
25
from jax.sharding import PartitionSpec as P
26

27
 import resources
28
from scaling_transformer_inference_efficiency import checkpoint
29
from scaling_transformer_inference_efficiency import partitioning
30

31

32
_TOY_HPARAMS = checkpoint.HParams(
33
    layers=3,
34
    embed=128,
35
    ff=256,
36
    heads=2,
37
    qkv=32,
38
    max_len=128,
39
    vocab=32128,
40
)
41

42

43
class PartitioningTest(absltest.TestCase):
44

45

46
  def test_copy_to_device_from_shape(self):
47
    shape = core.ShapedArray((4, 4), dtype=jnp.bfloat16)
48
    mesh = partitioning.make_mesh()
49
    x = partitioning.copy_to_device(shape,
50
                                    NamedSharding(mesh, P('x', ('y', 'z'))),
51
                                    shape)
52
    self.assertEqual(x.shape, shape.shape)
53
    self.assertEqual(x.dtype, shape.dtype)
54

55
  def test_copy_to_device_from_array(self):
56
    array = jnp.zeros((4, 4), jnp.bfloat16)
57
    shape = core.ShapedArray((4, 4), dtype=jnp.bfloat16)
58
    mesh = partitioning.make_mesh()
59
    x = partitioning.copy_to_device(array,
60
                                    NamedSharding(mesh, P('x', ('y', 'z'))),
61
                                    shape)
62
    self.assertEqual(x.shape, array.shape)
63
    self.assertEqual(x.dtype, array.dtype)
64

65

66
if __name__ == '__main__':
67
  absltest.main()
68