Numpyro: Crash / Infinite loop on second MCMC run

[ross-h1]

Hi guys... I'm experienceing a crash on the second run of a model. I have an application which needs to re fit a model several times as part of a backtesting strategy. The first iteration of the model runs through fine, but iteration two hangs or seems to get caught in an infinite loop. The issue seems to be in the _hashable() function in the mcmc.py file.

My (self contained) code with toy data attached. I have simpler model which has run over several iterations ok, but this more computationally expensive version fails...

Thanks in advance,

Ross

toy_model_PYRO.pdf

[fehiepsi]

Hi @ross-h1, I am not sure what we can do in this case. Your model has about 2000 latent variables, together with several scan loop. So it should be very slow. I made a quick check and found that

from numpyro.infer.util import initialize_model
from jax import value_and_grad
init_params1, potential_fn, _ = initialize_model(random.PRNGKey(0), tvmodel, model_kwargs=dict(X=y_jx,s=s_jx))
init_params2, _, _ = initialize_model(random.PRNGKey(0), tvmodel, model_kwargs=dict(X=y_jx,s=s_jx))
fn = value_and_grad(potential_fn)

%time jit(fn)(init_params1)[0].copy()
%time jit(fn)(init_params2)[0].copy()

1 leapfrog step would take about 8ms so typtically you can get a few samples in a minute (if chains mixing well). Could you try to call mcmc.print_summary() after each MCMC run to make sure that async dispatch does not affect the next run?

I would recommend using SVI or other methods to train this model. If your data is larger, you can consider running your model in GPU and vectorize garch, dcb computation. Basically, those functions compute the sequence x where

x[n] = a[n] + b * x[n-1] with x[-1] = x_init 

This can be vectorized by constructing a matrix of terms a and b. You can save some computation by computing terms a such as (1 - alpha - beta) * v0 + alpha * (x ** 2) or hxx0 * (1 - alpha - beta) + alpha * x ** 2, hxy0 * (1 - alpha - beta) + alpha * np.einsum('ab,ac->abc', y, x) outside scan loop.

[ross-h1]

Yep the print_summary() shows that the model
Is still crunching away, and hasn’t in fact finished in 107 seconds as originally suggested!

Am still struggling to get Numpyro working on GPU, will have another try tomorrow and try out the additional vectorisation you suggest.

Thank you, Ross

On 19 Jan 2020, at 05:08, Du Phan [e-mail] wrote:


Hi @ross-h1, I am not sure what we can do in this case. Your model has about 2000 latent variables, together with several scan loop. So it should be very slow. I made a quick check and found that

from numpyro.infer.util import initialize_model
from jax import value_and_grad
init_params1, potential_fn, _ = initialize_model(random.PRNGKey(0), tvmodel, model_kwargs=dict(X=y_jx,s=s_jx))
init_params2, _, _ = initialize_model(random.PRNGKey(0), tvmodel, model_kwargs=dict(X=y_jx,s=s_jx))
fn = value_and_grad(potential_fn)

%time jit(fn)(init_params1)[0].copy()
%time jit(fn)(init_params2)[0].copy()
1 leapfrog step would take about 8ms so typtically you can get a few samples in a minute (if chains mixing well). Could you try to call mcmc.print_summary() after each MCMC run to make sure that async dispatch does not affect the next run?

I would recommend using SVI or other methods to train this model. If your data is larger, you can consider running your model in GPU and vectorize garch, dcb computation. Basically, those functions compute the sequence x where

x[n] = a[n] + b * x[n-1] with x[-1] = x_init
This can be vectorized by constructing a matrix of terms a and b. You can save some computation by computing terms a such as (1 - alpha - beta) * v0 + alpha * (x * 2) or hxx0 * (1 - alpha - beta) + alpha * x * 2, hxy0 * (1 - alpha - beta) + alpha * np.einsum('ab,ac->abc', y, x) outside scan loop.

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub, or unsubscribe.

[fehiepsi]

@ross-h1 I implemented the solve for the sequence x[n] = a[n] + b * x[n-1]. You can use it in both garch and dcb. I just tested it with your data, it has the same speed as your CPU non-vectorized version so it might not worth to be implemented. Unfortunately, the MCMC result is pretty bad. :(

def inhomogeneous_scan(x_init, a, b):
    """
    :param x: a vector with shape (...,)
    :param a: a vector with shape (n, ...)
    :param b: a scalar

    Returns solution of the system x[n] = a[n] + b * x[n-1] with x[-1] = x_init 
    """
    x_shape = x_init.shape
    x_init = x_init.reshape(-1)
    n = a.shape[0]
    a = a.reshape((n, -1)).T
    # construct b_matrix
    power = np.expand_dims(np.arange(1, n + 1), 1) - np.arange(0, n + 1)
    b_matrix = np.tril(np.power(b, power), k=1)
    # construct a_matrix
    a_matrix = np.pad(np.broadcast_to(np.expand_dims(a, 1), a.shape + (n,)),
                      ((0, 0), (0, 0), (1, 0)), constant_values=1)
    # contruct the linear system
    A = b_matrix * a_matrix
    x = np.concatenate([np.expand_dims(x_init, 1), np.ones(a.shape)], -1)[..., None]
    res = (A @ x).squeeze(-1).T
    return res.reshape((n,) + x_shape)

[ross-h1]

Thank you, I built something similar, and got to more or less the same place. I could get a function to calculate my logp running in JAX, but it was about as fast on CPU and GPU. I still cant get Numpyro to work on at all on GPU, even with a simpler non time varying regression model, it throws an XLA compile error… I have a 2 GPU machine that I hoped to be able to use to run two chains in parallel… Will have another go tomorrow. I’m wondering if I might have to try another approach for this model, thinking about sequential / particle based approach, GPU doesn’t really speed up GARCH looping very much!

Much appreciate your help / efforts, thank you.

On 22 Jan 2020, at 20:07, Du Phan [e-mail] wrote:

@ross-h1 https://github.com/ross-h1 I implemented the solve for the sequence x[n] = a[n] + b * x[n-1]. You can use it in both garch and dcb. I just tested it with your data, it has the same speed as your CPU non-vectorized version so it might not worth to be implemented. Unfortunately, the MCMC result is pretty bad. :(

def inhomogeneous_scan(x_init, a, b):
"""
:param x: a vector with shape (...,)
:param a: a vector with shape (n, ...)
:param b: a scalar

Returns solution of the system x[n] = a[n] + b * x[n-1] with x[-1] = x_init 
"""
x_shape = x_init.shape
x_init = x_init.reshape(-1)
n = a.shape[0]
a = a.reshape((n, -1)).T
# construct b_matrix
power = np.expand_dims(np.arange(1, n + 1), 1) - np.arange(0, n + 1)
b_matrix = np.tril(np.power(b, power), k=1)
# construct a_matrix
a_matrix = np.pad(np.broadcast_to(np.expand_dims(a, 1), a.shape + (n,)),
                  ((0, 0), (0, 0), (1, 0)), constant_values=1)
# contruct the linear system
A = b_matrix * a_matrix
x = np.concatenate([np.expand_dims(x_init, 1), np.ones(a.shape)], -1)[..., None]
res = (A @ x).squeeze(-1).T
return res.reshape((n,) + x_shape)

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub https://github.com/pyro-ppl/numpyro/issues/528?email_source=notifications&email_token=AF5RD3OK6QK4VS5QE7363Z3Q7CRPFA5CNFSM4KIR7WRKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEJU5V7I#issuecomment-577362685, or unsubscribe https://github.com/notifications/unsubscribe-auth/AF5RD3LK3NR4GZNEGGXLSBDQ7CRPFANCNFSM4KIR7WRA.

[fehiepsi]

@ross-h1 We just made a new release, which you can try to see if GPU works for you. If you can use GPU with JAX codes but not with NumPyro code, please let us know (this indicates that there is something wrong with our configurations).

As I understood, the topic issue has been identified. So I close this for now.

[ross-h1]

Hi there, thank you! Now running on 2x GPU. I have switched to a different model for dynamic variances and covariances which eliminates the GARCH iteration, so can all be vectorised. Finally seeing some benefit of spending money with Nvidia... Ross

On 25 Jan 2020, at 23:04, Du Phan [e-mail] wrote:


@ross-h1 We just made a new release, which you can try to see if GPU works for you. If you can use GPU with JAX codes but not with NumPyro code, please let us know (this indicates that there is something wrong with our configurations).

As I understood, the topic issue has been identified. So I close this for now.

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub, or unsubscribe.