Hi,

Thanks for creating this amazing library!

So from what I understood this is the minimal needed to benefit from the JMP speedup

policy = jmp.Policy(param_dtype=jnp.float32, compute_dtype=jnp.float16, output_dtype=jnp.float16)
...
# creating a network, and creating a loss_fn using the network
data = policy.cast_to_compute(data)
params = policy.cast_to_compute(params)

grads = jax.grad(loss_fn)(data, params)
...
grads = policy.cast_to_param(grads)

The loss scale is only needed if we experience NaN or inf values. And thus we should see a difference in the time needed for a jax.grad operation when the data and params are in float16.

Please correct me if I'm wrong or if I miss anything.

Considering that, I have some questions:

• Can we expect to see a 2x or any speedup if we time the jax.grad operation with (params, data) in float16 and in float32 or is the speedup only achieved at scale ?
• Can we expect to see a 2x or any speedup with small networks and small experiments ? (say 2-layer nets and experiments that only need 2-3 minutes on a single gpu)
• Can we expect to see a 2x or any speedup with all types of networks or only some with specific architectures ? (say a 50 layer MLP, Transformers, etc)
• Is it necessary to apply the mixed precision policy to the network and thus to use Haiku for hk.mixed_precision.set_policy or is having the parameters and data in float16 sufficient to have a speedup even with a network created by us ?

Thank you and have a nice day !

Hi,

Thanks for creating this amazing functionality!

I have a basic question about the use of policy functions to set certain precision levels. As stated in the example of your README.md:

def layer(params, x):
  params, x = my_policy.cast_to_compute((params, x))
  w, b = params
  y = x @ w + b
  return my_policy.cast_to_output(y)

I am new to JAX, but the first thing I learned is that JAX likes pure functions. Does the use of my_policy violate the pure functions paradigm?

Should it become:

def layer(params, x, my_policy):
  params, x = my_policy.cast_to_compute((params, x))
  w, b = params
  y = x @ w + b
  return my_policy.cast_to_output(y)

The function can be jitted by using partial.

from functools import partial
layer_compiled = jit(partial(layer, my_policy=my_policy))

Fundamental question

A more fundamental question (which I maybe need to ask at JAX ) is: How should functions as input to functions be handled in JAX?

func_a(x,w,func):
  ...

func_a_compiled = jit(partial(func_a, func=func_b))

In order to jit this, I came up with the partial solution above. I assume your use of my_policy is valid, as you probably have more experience with JAX. But that creates some magic, which is undesirable for my use case. Is the jit(partial()) solution valid or is there a better way to handle functions as input to functions?

Have a nice day! J

When trying to build my own wheel from the tar-ball of jmp-0.0.2 that has been uploaded to PyPI I get the following error:

Collecting jmp==0.0.2
  Downloading jmp-0.0.2.tar.gz (13 kB)
    Running command python setup.py egg_info
    Traceback (most recent call last):
      File "<string>", line 1, in <module>
      File "/tmp/pip-download-zowwz99o/jmp_55ae4a2c4cdb44aa99c65fbf7a5ee9bb/setup.py", line 55, in <module>
        install_requires=_parse_requirements('requirements.txt'),
      File "/tmp/pip-download-zowwz99o/jmp_55ae4a2c4cdb44aa99c65fbf7a5ee9bb/setup.py", line 32, in _parse_requirements
        with open(requirements_txt_path) as fp:
    FileNotFoundError: [Errno 2] No such file or directory: 'requirements.txt'
WARNING: Discarding https://files.pythonhosted.org/packages/7c/ba/a6bfcaeedca8551e2fb4054d1fd061a0dd97d26dd44002b3e92d13b51877/jmp-0.0.2.tar.gz#sha256=fdb5cec0d10aab4116c2770f24b2adf4f503fcfbb96ce8ef583e1879bdbf1b9b (from https://pypi.org/simple/jmp/). Command errored out with exit status 1: python setup.py egg_info Check the logs for full command output.
ERROR: Could not find a version that satisfies the requirement jmp==0.0.2 (from versions: 0.0.1, 0.0.2)
ERROR: No matching distribution found for jmp==0.0.2

After downloading and extracting jmp-0.0.2.tar.gz from PyPI and looking at setup.py, I see that it contains lines that reference requirements.txt:

install_requires=_parse_requirements('requirements.txt')

however that file is not part of the source distribution and my build fails.

Oliver

P.S. Yes I know that there is a wheel for jmp-0.0.2 on PyPI, which I ended up using, but I'm using our Wheels_builder script that will recursively build wheels for dependencies for our systems. My point is that the source distribution of jmp 0.0.2 is incomplete as it lacks the information that is contained in the requirements file.

I'll first start by saying that I've just starting with JAX, so I might be doing something wrong.

When I run the following code:

precision = jmp.get_policy('params=float32,compute=float16,output=float16')
some_input = np.arange(15).reshape((5, 3))
@jax.jit
def some_function(some_input):
    some_input = precision.cast_to_compute(some_input)
    print(some_input.dtype)
    # prints float32
   return a_float16_compute_model(some_input) # fails

It seems that at least on the first (tracing) run, the numpy array doesn't get cast to float16, maybe because it is being treated as a tree, here

https://github.com/deepmind/jmp/blob/4b94370b8de29b79d6f840b09d1990b91c1afddd/jmp/_src/policy.py#L26

Surprisingly, if I run

some_input = np.arange(15).reshape((5, 3)).astype(precision.compute_dtype)
print(some_input.dtype)

#prints float16

the cast succeeds. I think that the expected behavior is that precision.cast_to_compute(some_input) should return a numpy array with compute_dtype, but I might be missing something.

[JAX] Fix test failure due to upcoming change to JAX.

An upcoming change to JAX adds a @jit decorator around a number of array operators, in this case division (/). A side effect of that change is that large Python integer constants that overflow an int32 or int64 type may produce an error. The workaround is either to explicitly cast the large constants to a specific type (e.g., np.float64), or in this case we can just do the math in question in classic NumPy since it is computing test expectations.

Bump version to 0.0.3.dev.

This is so if users install from GitHub we can see they are not using a stable version in bug reports:

$ pip install git+https://github.com/deepmind/jmp
$ python -c 'import jmp ; print(jmp.__version__)'
0.0.3.dev

Cut 0.0.2 release and create GitHub action to publish to PyPi.

We are not using 0.0.1 because this version was already used by the previous owner of the pypi package (https://pypi.org/project/jmp/0.0.1/). As such our releases will start at 0.0.2.

Thanks for the awesome work!

This PR fixes an issue where jax.lax.select complains about dtypes not being equal when adjusting the DynamicLossScale.

The exception's stack trace ends with:

File ".../jmp/_src/loss_scale.py", line 147, in adjust
    loss_scale = jax.lax.select(
TypeError: lax.select requires arguments to have the same dtypes, got float32, int32. (Tip: jnp.where is a similar function that does automatic type promotion on inputs).

My code looks similar to

scale = jmp.DynamicLossScale(jnp.asarray(2 ** 15))
...
gradients, scale = gradient_fn(..., scale)
gradients = scale.unscale(gradients)
gradients_finite = jmp.all_finite(gradients)
scale = scale.adjust(gradients_finite)  # This line throws the exception
...