Large Performance Gap Between Jax and Pytorch on Tree Traversal

[jamespinkerton]

Hi. I'm writing code to do inference on decision tree models. The ideas come from https://github.com/microsoft/hummingbird.

I'm getting 4x or so slower performance in my jax implementation and I can't figure out if it's my fault or inherently slower in jax. I'm running this on a CPU.

I'm guessing this is me writing something very slow in jax and that it's my fault. Any help would be very appreciated!

  import time
  from typing import Any, NamedTuple

  import jax
  import jax.numpy as jnp
  import lightgbm as lgb  # type: ignore
  import numpy as np
  import torch


  class TreeParameters(NamedTuple):
      lefts: list[int]
      rights: list[int]
      features: list[int]
      thresholds: list[float]
      values: list[float]
      depth: int


  class TreeTraversalTreeImpl(torch.nn.Module):
      def __init__(self, tree_parameters: list[TreeParameters]):
          super().__init__()
          self.max_tree_depth = max([tp.depth for tp in tree_parameters])
          self.num_trees = len(tree_parameters)
          self.num_nodes = max([len(tp.features) for tp in tree_parameters])

          lefts = torch.zeros((self.num_trees, self.num_nodes), dtype=torch.int64)
          rights = torch.zeros((self.num_trees, self.num_nodes), dtype=torch.int64)
          features = torch.zeros((self.num_trees, self.num_nodes), dtype=torch.int64)
          thresholds = torch.zeros((self.num_trees, self.num_nodes))
          values = torch.zeros((self.num_trees, self.num_nodes))

          for i, tp in enumerate(tree_parameters):
              s = len(tp.features)
              lefts_list = [i if v == -1 else v for i, v in enumerate(tp.lefts)]
              rights_list = [i if v == -1 else v for i, v in enumerate(tp.rights)]
              lefts[i][:s] = torch.tensor(lefts_list, dtype=torch.int64)
              rights[i][:s] = torch.tensor(rights_list, dtype=torch.int64)
              features[i][:s] = torch.tensor(tp.features, dtype=torch.int64)
              thresholds[i][:s] = torch.tensor(tp.thresholds)
              values[i][:s] = torch.tensor(tp.values)

          def detached_param(x: torch.Tensor) -> torch.Tensor:
              return torch.nn.Parameter(x.view(-1).detach().clone(), requires_grad=False)

          self.lefts = detached_param(lefts)
          self.rights = detached_param(rights)
          self.features = detached_param(features)
          self.thresholds = detached_param(thresholds)
          self.values = detached_param(values)

          nodes_offset = [[i * self.num_nodes for i in range(self.num_trees)]]
          self.nodes_offset = torch.nn.Parameter(
              torch.LongTensor(nodes_offset), requires_grad=False
          )

      def select(self, x: torch.Tensor, indexes: torch.Tensor) -> torch.Tensor:
          return torch.index_select(x, 0, indexes).view(-1, self.num_trees)

      def forward(self, x: torch.Tensor) -> torch.Tensor:
          indexes = self.nodes_offset.expand(x.size()[0], self.num_trees).reshape(-1)

          for i in range(self.max_tree_depth):
              thresholds = self.select(self.thresholds, indexes)
              lefts = self.select(self.lefts, indexes)
              rights = self.select(self.rights, indexes)
              feature_nodes = self.select(self.features, indexes)
              feature_values = torch.gather(x, 1, feature_nodes)
              indexes = torch.where(feature_values <= thresholds, lefts, rights).long()
              indexes = indexes + self.nodes_offset
              indexes = indexes.view(-1)

          return self.select(self.values, indexes).sum(1)

  class TreeTraversalJax:
      def __init__(self, tree_parameters: list[TreeParameters]):
          self.max_tree_depth = max([tp.depth for tp in tree_parameters])
          self.num_trees = len(tree_parameters)
          self.num_nodes = max([len(tp.features) for tp in tree_parameters])

          lefts = np.zeros((self.num_trees, self.num_nodes), dtype=np.int32)
          rights = np.zeros((self.num_trees, self.num_nodes), dtype=np.int32)
          features = np.zeros((self.num_trees, self.num_nodes), dtype=np.int32)
          thresholds = np.zeros((self.num_trees, self.num_nodes))
          values = np.zeros((self.num_trees, self.num_nodes))

          for i, tp in enumerate(tree_parameters):
              s = len(tp.features)
              lefts_list = [i if v == -1 else v for i, v in enumerate(tp.lefts)]
              rights_list = [i if v == -1 else v for i, v in enumerate(tp.rights)]
              lefts[i][:s] = np.array(lefts_list, dtype=np.int32)
              rights[i][:s] = np.array(rights_list, dtype=np.int32)
              features[i][:s] = np.array(tp.features, dtype=np.int32)
              thresholds[i][:s] = np.array(tp.thresholds)
              values[i][:s] = np.array(tp.values)

          self.lefts = jnp.array(lefts).reshape(-1)
          self.rights = jnp.array(rights).reshape(-1)
          self.features = jnp.array(features).reshape(-1)
          self.thresholds = jnp.array(thresholds).reshape(-1)
          self.values = jnp.array(values).reshape(-1)

          nodes_offset = [[i * self.num_nodes for i in range(self.num_trees)]]
          self.nodes_offset = jnp.array(nodes_offset, dtype=jnp.int32)

      def select(self, x: jax.Array, indexes: jax.Array) -> jax.Array:
          return jnp.take(x, indexes, axis=0).reshape(-1, self.num_trees)

      def __call__(self, x: jax.Array) -> jax.Array:
          indexes = jnp.broadcast_to(
              self.nodes_offset, (x.shape[0], self.num_trees)
          ).reshape(-1)

          for i in range(self.max_tree_depth):
              thresholds = self.select(self.thresholds, indexes)
              lefts = self.select(self.lefts, indexes)
              rights = self.select(self.rights, indexes)
              feature_nodes = self.select(self.features, indexes)
              feature_values = jnp.take_along_axis(x, feature_nodes, axis=1)
              indexes = jnp.where(feature_values <= thresholds, lefts, rights).astype(
                  jnp.int32
              )
              indexes = indexes + self.nodes_offset
              indexes = indexes.reshape(-1)

          return self.select(self.values, indexes).sum(1)

  def _lgbm_tree_traversal(
      node: Any,
      lefts: list[int],
      rights: list[int],
      features: list[int],
      thresholds: list[float],
      values: list[float],
      count: int,
      depth: int,
  ) -> tuple[int, int]:
      if "left_child" in node:
          features.append(node["split_feature"])
          thresholds.append(node["threshold"])
          values.append(-1)
          lefts.append(count + 1)
          rights.append(-1)
          pos = len(rights) - 1
          count, depth_left = _lgbm_tree_traversal(
              node["left_child"],
              lefts,
              rights,
              features,
              thresholds,
              values,
              count + 1,
              depth + 1,
          )
          rights[pos] = count + 1
          count, depth_right = _lgbm_tree_traversal(
              node["right_child"],
              lefts,
              rights,
              features,
              thresholds,
              values,
              count + 1,
              depth + 1,
          )
          return count, max(depth_left, depth_right)
      else:
          features.append(0)
          thresholds.append(0)
          values.append(node["leaf_value"])
          lefts.append(-1)
          rights.append(-1)
          return count, depth


  def _parse_lgbm_tree_parameters(tree_info: Any) -> TreeParameters:
      lefts: list[int] = []
      rights: list[int] = []
      features: list[int] = []
      thresholds: list[float] = []
      values: list[float] = []
      count, depth = _lgbm_tree_traversal(
          tree_info["tree_structure"], lefts, rights, features, thresholds, values, 0, 0
      )
      return TreeParameters(
          lefts=lefts,
          rights=rights,
          features=features,
          thresholds=thresholds,
          values=values,
          depth=depth,
      )

  def parse_lgbm_model(model: Any) -> list[TreeParameters]:
      trees = model.dump_model()["tree_info"]
      return [_parse_lgbm_tree_parameters(t) for t in trees]


  if __name__ == "__main__":
      np.random.seed(0)

      max_depth = 10
      num_leaves = 63
      n_feautures = 100
      num_trees = 400
      batch = 200
      S = 10_000
      X_train = np.random.randn(S, n_feautures)
      X_train_torch = torch.from_numpy(X_train).float()
      X_train_jax = jnp.array(X_train, dtype=jnp.float32)
      y_train = np.random.randn(S)

      train_data = lgb.Dataset(X_train, label=y_train)
      model = lgb.train(
          {"max_depth": max_depth, "num_leaves": num_leaves},
          train_data,
          num_boost_round=num_trees,
      )
      serialized_model = parse_lgbm_model(model)
      tree_model = TreeTraversalTreeImpl(serialized_model)
      tree_model = torch.jit.freeze(torch.jit.script(tree_model.eval()))
      tree_model_jax_jit = jax.jit(TreeTraversalJax(serialized_model))

      count = len(X_train) // batch

      r = tree_model_jax_jit(X_train_jax)
      print(r)
      r = tree_model(X_train_torch)
      print(r)

      begin = time.time()
      for i in range(count):
          b1 = time.time()
          tree_model_jax_jit(X_train_jax[i * batch : (i + 1) * batch])
          e1 = time.time()
          print("jax", (e1 - b1) * 1e3)
      end = time.time()
      print("average jax", (end - begin) / count * 1e3)

      begin = time.time()
      for i in range(count):
          b1 = time.time()
          tree_model(X_train_torch[i * batch : (i + 1) * batch])
          e1 = time.time()
          print("torch", (e1 - b1) * 1e3)
      end = time.time()
      print("average torch", (end - begin) / count * 1e3)

[jakevdp]

On the JAX side, I suspect what you're running into is long compile times. Under JIT, JAX will unroll all Python loops into a linear program and pass this flat set of instructions to XLA, which can result in large programs with long compile times.

I'd suggest one of the following:

• if the interior of the loop is relatively expensive, you might consider JIT-compiling this interior code, and running the for loop outside the JIT context. Then you'd get the benefit of JIT compilation within each iteration, but not incur the cost of compiling the flattened program.
• if the interior of the loop is relatively inexpensive, you might be better off replacing your Python for loop with lax.fori_loop. See JAX Sharp Bits: structured control flow for more information about this.

Finally, I should note that your JAX benchmarks may not be accurate because of JAX's asynchronous dispatch; see FAQ: Benchmarking JAX Code for some tips on getting accurate benchmarks of JAX functions. In particular, you probably want to wrap your benchmark code in jax.block_until_ready().

Good luck!

[jakevdp]

Your benchmark includes compilation time, which is why I mentioned it.

[jamespinkerton]

I'm pretty sure the benchmark doesn't include compilation time. Why do you think it does?

[jakevdp]

A function wrapped in jax.jit will be compiled the first time it is called with inputs of a particular shape and dtype (this is the "Just In Time" part of JIT compilation). The first time you call tree_model_jax_jit with an input of shape (200, 100) is within the benchmark loop, and so the compilation time will be part of your benchmark measurement.

[jakevdp]

You can move JAX's compilation outside the benchmark loop by putting something like this before your benchmark loop:

begin = time.time()
_ = tree_model_jax_jit(X_train_jax[0: batch])
end = time.time()
print("JAX compilation time", (end - begin) * 1e3)

When I add that to your code, I get this output:

JAX compilation time 5440.757751464844
...
average jax 50.66882610321045
...
average torch 58.736515045166016

You see that once you exclude compilation time from the benchmark, JAX is actually faster than pytorch (this is on a standard Colab CPU runtime).

Now, compilation time is a real cost that you have to pay sometime, and if you want to address that I'd suggest the strategies in my answer above.

[jamespinkerton]

I have two print statements in my benchmark. I wasn't referring to the "average" line in my benchmark! I was referring to the last non-average line of each. And so for those lines it is not the first time it's being called. So when we're citing 7ms or 24ms (without your lax suggestion), we're not talking about the average that includes compilation time. But it looks like you're right that there is re-compilation happening every time that jax is calling the function when a regular for loop is used, and I'm not only paying for compilation on the first call.

[JiaYaobo]

@jamespinkerton just like what @jakevdp mentioned, you can use lax.fori_loop for accelerating your code! here's a small modification for your __call__ method.

@partial(jax.jit, static_argnums=(0,))
def __call__(self, x: jax.Array) -> jax.Array:
    indexes = jnp.broadcast_to(
        self.nodes_offset, (x.shape[0], self.num_trees)
    ).reshape(-1)

    def fori_loop_body(i, indexes):
        thresholds = self.select(self.thresholds, indexes)
        lefts = self.select(self.lefts, indexes)
        rights = self.select(self.rights, indexes)
        feature_nodes = self.select(self.features, indexes)
        feature_values = jnp.take_along_axis(x, feature_nodes, axis=1)
        indexes = jnp.where(feature_values <= thresholds, lefts, rights).astype(
            jnp.int32
        )
        indexes = indexes + self.nodes_offset
        indexes = indexes.reshape(-1)
        return indexes
    indexes = jax.lax.fori_loop(0, self.max_tree_depth, fori_loop_body, indexes)
    return self.select(self.values, indexes).sum(1)

As a result, average torch 7.899 and average jax 7.568, and jax spends about 5 ms for each loop except first one on my m1 pro chip macbook

[jamespinkerton]

That does help a lot! I Thank you to both of you. I have two questions:

1. Jake said "if the interior of the loop is relatively inexpensive, you might be better off replacing your Python for loop with lax.fori_loop". I was under the impression that the interior of this for loop is relatively expensive, as it's almost the only work done in the whole function! Why does this answer work anyway?
2. On my google cloud C3 machine I get pytorch at 9ms and jax at 9.93ms with your this implementation. Running on my mac I get 3.7ms for torch and 3.3ms for jax! It's hard for me to understand the big differences here. Any suggestions for how I'm supposed to understand how to get parity on the cloud box?

[jakevdp]

I was under the impression that the interior of this for loop is relatively expensive, as it's almost the only work done in the whole function! Why does this answer work anyway?

I answered this at the top of my answer:

Under JIT, JAX will unroll all Python loops into a linear program and pass this flat set of instructions to XLA, which can result in large programs with long compile times.

By contrast, if you're using fori_loop, you only compile a single copy of the loop body and XLA repeatedly applies this compiled body to the inputs.