The JIT time of my model is too long. Is there anything that can be improved?

[nasyxx]

I'm not quite sure what I need to present here.

The JAX version is 0.4.16, the Flax version is 0.7.4, H100 and CUDA12 are used, and CUDA_VISIBLE_DEVICES is also used to limit seeing only 1 GPU.

Here is the model, it's a very simple UNet.

from collections.abc import Sequence
from dataclasses import field

import jax
import jax.numpy as jnp
from flax import linen

PAD_ = str | int | Sequence[int | tuple[int, int]]


class DoubleConv(linen.Module):
  """Double Convolution Layer."""

  out_channels: int
  mid_channel: int | None = None
  kernel_size: int = 3
  stride: int = 1
  padding: PAD_ = "SAME"
  with_bn: bool = True
  bn_config: dict = field(
      default_factory=lambda: {
          "momentum": 0.99,
          "epsilon": 1e-4,
          "use_scale": True,
          "use_bias": True,
      }
  )

  @linen.compact
  def __call__(self, x: jax.Array, train: bool = False) -> jax.Array:
    """Call."""
    mid_channel = self.mid_channel or self.out_channels

    x = linen.Conv(
        mid_channel,
        (self.kernel_size, self.kernel_size),
        self.stride,
        self.padding,
    )(x)
    if self.with_bn:
      x = linen.BatchNorm(**self.bn_config)(x, use_running_average=not train)
    x = jax.nn.relu(x)

    x = linen.Conv(
        self.out_channels,
        (self.kernel_size, self.kernel_size),
        self.stride,
        self.padding,
    )(x)
    if self.with_bn:
      x = linen.BatchNorm(**self.bn_config)(x, use_running_average=not train)
    return jax.nn.relu(x)


class Down(linen.Module):
  """Down Sampling Layer."""

  out_channels: int

  @linen.compact
  def __call__(self, x: jax.Array, train: bool = True) -> jax.numpy.ndarray:
    """Call."""
    x = linen.max_pool(x, window_shape=(2, 2), strides=(2, 2), padding="VALID")
    return DoubleConv(out_channels=self.out_channels)(x, train=train)


class UpSample(linen.Module):
  """Up Sampling Layer."""

  out_channels: int
  kernel_size: int = 2
  stride: int = 2
  padding: PAD_ = "SAME"

  @linen.compact
  def __call__(self, x1: jax.Array, x2: jax.Array, train: bool = True) -> jax.Array:
    """Call."""
    x1 = linen.ConvTranspose(
        self.out_channels,
        (self.kernel_size, self.kernel_size),
        (self.stride, self.stride),
        self.padding,
    )(x1)
    diff1 = x2.shape[1] - x1.shape[1]
    diff2 = x2.shape[2] - x1.shape[2]
    x1 = jnp.pad(
        x1,
        [
            (0, 0),
            (diff1 // 2, diff1 - diff1 // 2),
            (diff2 // 2, diff2 - diff2 // 2),
            (0, 0),
        ],
    )
    x = jnp.concatenate([x1, x2], axis=-1)
    return DoubleConv(out_channels=self.out_channels)(x, train=train)

class UNet(linen.Module):
  """UNet."""

  out_channels: int
  kernel_size: int = 3
  stride: int = 1
  padding: str = "SAME"
  bn_config: dict = field(
      default_factory=lambda: {
          "momentum": 0.99,
          "epsilon": 1e-4,
          "use_scale": True,
          "use_bias": True,
      }
  )

  @linen.compact
  def __call__(self, x: jax.Array, train: bool = False) -> jax.numpy.ndarray:
    """Call."""
    x1 = DoubleConv(  # type: ignore
        out_channels=self.out_channels,
        kernel_size=self.kernel_size,
        stride=self.stride,
        padding=self.padding,
        bn_config=self.bn_config,
    )(x, train=train)
    x2 = Down(out_channels=self.out_channels * 2)(x1, train=train)
    x3 = Down(out_channels=self.out_channels * 4)(x2, train=train)
    x4 = Down(out_channels=self.out_channels * 8)(x3, train=train)
    x5 = Down(out_channels=self.out_channels * 16)(x4, train=train)

    x = UpSample(out_channels=self.out_channels * 8)(x5, x4, train=train)
    x = UpSample(out_channels=self.out_channels * 4)(x, x3, train=train)
    x = UpSample(out_channels=self.out_channels * 2)(x, x2, train=train)
    x = UpSample(out_channels=self.out_channels)(x, x1, train=train)

    return linen.Conv(
        features=1,
        kernel_size=(self.kernel_size, self.kernel_size),
        strides=self.stride,
        padding=self.padding,
    )(x)

Here is the log

2023-09-29 15:37:06.725613: W external/tsl/tsl/framework/bfc_allocator.cc:296] Allocator (GPU_0_bfc) ran out of memory trying to allocate 65.08GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
2023-09-29 15:37:07.020458: W external/tsl/tsl/framework/bfc_allocator.cc:296] Allocator (GPU_0_bfc) ran out of memory trying to allocate 65.08GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
2023-09-29 15:37:10.779051: W external/tsl/tsl/framework/bfc_allocator.cc:296] Allocator (GPU_0_bfc) ran out of memory trying to allocate 64.63GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
2023-09-29 15:37:11.359490: W external/tsl/tsl/framework/bfc_allocator.cc:296] Allocator (GPU_0_bfc) ran out of memory trying to allocate 64.63GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
2023-09-29 15:37:21.426731: E external/xla/xla/service/slow_operation_alarm.cc:65] Trying algorithm eng35{k2=1,k5=2,k14=5} for conv (f32[64,128,2,2]{3,2,1,0}, u8[0]{0}) custom-call(f32[64,2,1000,1000]{3,2,1,0}, f32[128,2,999,999]{3,2,1,0}), window={size=999x999 rhs_reversal=1x1}, dim_labels=bf01_oi01->bf01, custom_call_target="__cudnn$convForward", backend_config={"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0} is taking a while...
2023-09-29 15:37:22.421563: E external/xla/xla/service/slow_operation_alarm.cc:133] The operation took 1.995028884s
Trying algorithm eng35{k2=1,k5=2,k14=5} for conv (f32[64,128,2,2]{3,2,1,0}, u8[0]{0}) custom-call(f32[64,2,1000,1000]{3,2,1,0}, f32[128,2,999,999]{3,2,1,0}), window={size=999x999 rhs_reversal=1x1}, dim_labels=bf01_oi01->bf01, custom_call_target="__cudnn$convForward", backend_config={"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0} is taking a while...
2023-09-29 15:37:23.867071: E external/xla/xla/service/slow_operation_alarm.cc:65] Trying algorithm eng38{k2=1,k13=0,k14=4,k18=1,k22=0,k23=0} for conv (f32[64,128,2,2]{3,2,1,0}, u8[0]{0}) custom-call(f32[64,2,1000,1000]{3,2,1,0}, f32[128,2,999,999]{3,2,1,0}), window={size=999x999 rhs_reversal=1x1}, dim_labels=bf01_oi01->bf01, custom_call_target="__cudnn$convForward", backend_config={"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0} is taking a while...
2023-09-29 15:37:24.613088: E external/xla/xla/service/slow_operation_alarm.cc:133] The operation took 1.746117634s
Trying algorithm eng38{k2=1,k13=0,k14=4,k18=1,k22=0,k23=0} for conv (f32[64,128,2,2]{3,2,1,0}, u8[0]{0}) custom-call(f32[64,2,1000,1000]{3,2,1,0}, f32[128,2,999,999]{3,2,1,0}), window={size=999x999 rhs_reversal=1x1}, dim_labels=bf01_oi01->bf01, custom_call_target="__cudnn$convForward", backend_config={"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0} is taking a while...
2023-09-29 15:37:25.613274: E external/xla/xla/service/slow_operation_alarm.cc:65] Trying algorithm eng39{k2=0,k12=-1,k13=2,k14=3,k15=0,k17=2,k18=1,k23=0} for conv (f32[64,128,2,2]{3,2,1,0}, u8[0]{0}) custom-call(f32[64,2,1000,1000]{3,2,1,0}, f32[128,2,999,999]{3,2,1,0}), window={size=999x999 rhs_reversal=1x1}, dim_labels=bf01_oi01->bf01, custom_call_target="__cudnn$convForward", backend_config={"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0} is taking a while...

[chiamp]

Can you provide a small reproducible example of the jitted training loop? How long does it take to compile?

[nasyxx]

Hi, sorry for responding late.

I could not reproduce on a small piece. And, I may find that the problem occurs in the deterministic of CNN. I tried adding XLA_FLAGS='--xla_gpu_deterministic_ops=true' and this error no longer occurs.

[nasyxx]

@chiamp

Hi, thanks for your response.

Can you provide a small reproducible example of the jitted training loop? How long does it take to compile?

To reproduce:

Here is the repo: https://github.com/uta-smile/lab_challenge_23fall

Install pdm from https://pdm.fming.dev/latest/

by

curl -sSL https://pdm.fming.dev/install-pdm.py | python3 -

And install all dependence with

pdm install ":all"

After that, we can see two situations.

CUDA_VISIBLE_DEVICES=0 pdm run python -m train.train --no-infer --no-predict

Will generate the following errors:

XLA_FLAGS='--xla_gpu_deterministic_ops=true' CUDA_VISIBLE_DEVICES=0 pdm run python -m train.train --no-infer --no-predict

It will not generate any error.