GPU only OOM errors #12778

sabrinastronomy · 2022-10-12T20:47:13Z

sabrinastronomy
Oct 12, 2022

Hi there, I've been reading through some of the OOM discussion questions, but I haven't been able to find one that answers my issue. I am trying to minimize a chi-squared distribution (See chi_sq_jax below) using JAX's gradient descent capabilities. In essence, we hope to reconstruct clean data from noisy data. This setup is kind of GAN-esque but with different goals. With a 256 x 256 array (we hope to eventually move to 3D), we are able to properly reconstruct a 128 x 128 array (called field below) within 10 minutes on the CPU. This doesn't seem super realistic as we hope to eventually move to arrays of 256^3. I continuously get OOM errors when I try to run this same example on the GPU. Smaller arrays work on the GPU, i.e., up to ~128^2.

Why is the GPU running out of memory when using a JAX optimizer on these (from looking at other posts) relatively small arrays?

Initial guess of jax.optimize.minimize: 256 x 256 noise array
Output: 256 x 256 original image reconstructed
GPUs: NVIDIA P100 Pascal (16G HBM2 memory)
JAX version: '0.3.22', Python=3.10

Things I've tried:

Setting these variables:
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE']='false'
os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION']='.10'
Implementing the optimizer outside of JAX and only using the grad function.

Code below:

import jax
import jax.numpy as jnp
from jax.scipy.optimize import minimize as minimize_jax
import matplotlib.pyplot as plt
from jax import jit
import powerbox as pbox
import numpy as np
from jax.config import config
from memory_profiler import profile

config.update('jax_enable_x64', True)
import os

class GradDescent2D:
    def __init__(self, bias, pbox_field=True, side_length=256, dimensions=2, compile=True, vectorize=True, verbose=False):
        self.bias = bias
        self.pbox_field = pbox_field
        self.side_length = side_length
        self.dim = dimensions
        self.verbose = verbose
        pb_data = self.create_better_normal_field()
        clean_data = jnp.asarray(pb_data.delta_x())
        self.rms = jnp.std(clean_data)

        # Assume that the noise is 10% of the rms
        data = clean_data + self.create_normal_field(scale_num=0.1)
        self.data = jnp.asarray(data)
        # We also need the power spectrum box
        self.pspec_box = jnp.asarray(pb_data.power_array())
        self.N_diag = self.rms ** 2 * jnp.ones((self.side_length * self.side_length))

        # Generate a starting point field that is just 0.5 * rms white noise
        s_field = self.create_normal_field(scale_num=0.5) # starting guess
        self.s_field = jnp.asarray(s_field).flatten()

        if compile and not vectorize:
            self.jit_vec_minimize = jax.jit(self.differentiate_2D_func)
            self.opt_result = self.jit_vec_minimize(data)
        elif compile and vectorize:
            self.jit_vec_minimize = jax.jit(jax.vmap(self.differentiate_2D_func))
            self.opt_result = self.jit_vec_minimize(data)
        else:
            self.opt_result = self.differentiate_2D_func(data)

        true_answer = clean_data / self.bias
        noisy_true_answer = data / self.bias
        result = jnp.array(jnp.reshape(self.opt_result, (nside, nside)))
        self.check_field(result, "result")
        self.check_field(true_answer, "true_answer")
        self.check_field(noisy_true_answer, "noisy_true_answer")

    def create_normal_field(self, scale_num):
        np_version = np.random.normal(scale=scale_num * self.rms, size=(self.side_length, self.side_length))
        return jnp.asarray(np_version)

    def create_jax_normal_field(self):
        seed = jax.random.PRNGKey(0)
        return jax.random.normal(seed, shape=(self.side_length, self.side_length))

    def create_better_normal_field(self):
        pb = pbox.PowerBox(
            N=self.side_length,  # Number of grid-points in the box = the sidelengths * the dimension
            dim=self.dim,  # 2D box
            pk=lambda k: self.bias**2 * 0.1 * k ** -2.,  # The power-spectrum
            boxlength=1.0,  # Size of the box (sets the units of k in pk)
            seed=1010  # Set a seed to ensure the box looks the same every time (optional)
        )
        return pb

    def bias_field(self, bias, field):
        """:param bias, -- bias upon which to bias current field
        """
        if self.verbose:
            print("Your field has been biased.")
        return bias * field

    def fft_jax(self, field): 
        return jnp.fft.fftshift(jnp.fft.fftn(jnp.fft.ifftshift(field)))

    def independent_only_jax(self, box): 
        """
        This takes a box that is meant to be a Fourier space representation
        of a box that is purely real in configuration space and outputs
        only the independent elements. It does so separately for the real
        and imaginary parts because there are more independent real parts
        than there are imaginary parts.

        THIS ONLY WORKS IN 2D RIGHT NOW
        """
        N = box.shape[0]
        # First get the upper half plane minus the bits that aren't independent
        first_row = box[0, :-(N // 2 - 1)]
        sandwiched_rows = box[1:N // 2].flatten()
        origin_row = box[N // 2, :-(N // 2 - 1)]

        real_part = jnp.concatenate((jnp.real(first_row),
                                     jnp.real(sandwiched_rows),
                                     jnp.real(origin_row)))

        imag_part = jnp.concatenate((jnp.imag(first_row[1:-1]),
                                     jnp.imag(sandwiched_rows),
                                     jnp.imag(origin_row[1:-1])))

        return real_part, imag_part


    def chi_sq_jax(self, candidate_field): 
        """
        This is the chi^2, of the form
        s^t S^-1 s + [d - f(s, lambda)]^t N^-1 [d - f(s, lambda)]
        where s is the candidate field (that we are trying to fit for)
        S is the covariance matrix

        Inputs:
        --Array of real numbers for s (the candidate field)
        """

        # Convert the fields back to 2D arrays
        candidate_field = jnp.reshape(candidate_field, (self.side_length, self.side_length))
        # FT and get only the independent modes
        fourier_box = self.fft_jax(candidate_field)
        fourier_nums_real, fourier_nums_imag = self.independent_only_jax(fourier_box)
        # Do the same for the power spectrum box. Assigning an imaginary part to the
        # original power spectrum box is just a trick to store pspec in the right format
        # to be compared to the imaginary parts
        pspec_indep_nums_re, pspec_indep_nums_im = self.independent_only_jax(self.pspec_box + 1j * self.pspec_box)
        # Need to also remove the mean term because our fields are zero mean so both
        # the pspec and the real part will be zero at the origin
        fourier_nums_real = fourier_nums_real[:-1]
        pspec_indep_nums_re = pspec_indep_nums_re[:-1]

        # Work out the prior term
        real_prior = jnp.dot(fourier_nums_real ** 2, (2 / pspec_indep_nums_re))  # Half variance for real
        imag_prior = jnp.dot(fourier_nums_imag ** 2, (2 / pspec_indep_nums_im))  # Half variance for imag
        prior = real_prior + imag_prior

        # Now compute the likelihood term, assuming that the noise is diagonal in configuration space
        discrepancy = self.data - self.bias_field(candidate_field, self.bias)
        likelihood = jnp.dot(discrepancy.flatten() ** 2, 1. / self.N_diag)

        return likelihood  # + prior

    def differentiate_2D_func(self, data):
        func = self.chi_sq_jax
        opt_result = minimize_jax(func, x0=self.s_field, method='BFGS')
        return opt_result.x

 
if __name__ == "__main__":

    # FT and get only the independent modes
    nsides = [8] # this breaks when nsides~256
        start = time.time()
        GradDescent2D(bias=bias, side_length=nside, compile=True, vectorize=False)
        end = time.time()
        print("total time: ", end - start)

The memory device profiling using pprof for the 8 x 8 and 128 x 128 versions don't seem very illuminating to me. Here it is for 128 x 128.