Operation very slow to compile since jax 0.4.36

[btaba]

Description

Hi folks,

We've been having slow compilation issues since jax 0.4.36 with some of our JAX code. The slow compilation (i.e. takes O(hours) to run instead of <2 minutes) pops up for certain devices and for certain environments (with seemingly benign code changes). It has been quite tricky to pin down the issue over the last couple of months, and I haven't been able to create a smaller MRE than the following (on my local RTX 4090):

Install packages

pip install --upgrade jax[cuda] jaxlib 
pip install --upgrade mujoco
pip install --upgrade mujoco_mjx
pip install --upgrade brax

mujoco==3.2.7
mujoco-mjx==3.2.7
brax==0.12.1

Run this Python code:

import functools
from mujoco_playground import registry
from mujoco_playground import wrapper
from mujoco_playground.config import manipulation_params
from brax.training.agents.ppo import train as ppo
from brax.training.agents.ppo import networks as ppo_networks


env_name = 'PandaRobotiqPushCube'
env = registry.load(env_name)
env_cfg = registry.get_default_config(env_name)

ppo_params = manipulation_params.brax_ppo_config(env_name)
ppo_training_params = dict(ppo_params)
network_factory = ppo_networks.make_ppo_networks
if "network_factory" in ppo_params:
  del ppo_training_params["network_factory"]
  network_factory = functools.partial(
      ppo_networks.make_ppo_networks,
      **ppo_params.network_factory
  )

train_fn = functools.partial(
    ppo.train, **dict(ppo_training_params),
    network_factory=network_factory,
)
make_inference_fn, params, metrics = train_fn(
    environment=env,
    wrap_env_fn=wrapper.wrap_for_brax_training,
)

The corresponding XLA dump is attached.

I also reran the same script with num_evals=0 within train_fn, and the code runs fine (the slow compilation occurs somewhere here). I'm attaching both the working and non-working XLA dumps. We would really appreciate any help on this issue.

xla_dump_hanging_compilation.tar.gz
xla_dump_working.tar.gz

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.5.0
jaxlib: 0.5.0
numpy:  1.26.4
python: 3.12.3 (main, Sep 10 2024, 15:47:39) [GCC 13.2.0]
device info: NVIDIA GeForce RTX 4090-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='btaba.mtv.corp.google.com', release='6.10.11-1rodete2-amd64', version='#1 SMP PREEMPT_DYNAMIC Debian 6.10.11-1rodete2 (2024-10-16)', machine='x86_64')


$ nvidia-smi
Tue Jan 28 13:30:18 2025       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.216.01             Driver Version: 535.216.01   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce RTX 4090        Off | 00000000:61:00.0 Off |                  Off |
|  0%   50C    P2              22W / 450W |    397MiB / 24564MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A   2663746      C   .../.pyenv/versions/mjx-312/bin/python      390MiB |
+---------------------------------------------------------------------------------------+

[btaba]

Related to:

google/brax#569

google-deepmind/mujoco_playground#11

@erikfrey

[kevinzakka]

@mattjj Would appreciate your help here 🙏

[jreiffers]

HLO reproducer: https://gist.github.com/jreiffers/b6b8427ef64c00e688e79fd5af25b571

bazel run --config=cuda //xla/tools:run_hlo_module -c opt  -- --xla_disable_all_hlo_passes --input_format=hlo --random_init_input_literals --platform=CUDA /tmp/repro.hlo

I'll take a look which pass is blowing up. Hopefully it's not in LLVM :-).

[jreiffers]

I initially suspected the unnecessary concatenate at the end, but it happens even when that's fixed manually. There are at least two issues here:

1. codegen/emitters/computation_partitioner.cc does a really poor job partitioning this, because of the bitcasts above the concatenate. The partitioner does not track the real indexing, so after the bitcast it considers the inputs to the bitcasts (e.g. add.56926.1.clone.1) to have users with inconsistent indexing (the root tuple and the bitcast), which cascades all the way to the parameters. We end up with lots of functions even though ~everything is elementwise.
2. Afterwards, inlining breaks, generating a huge function. I stopped it after a few hundred thousand ops. Not entirely sure yet what's happening here, maybe failed/insufficient CSE or canonicalization.

The correct fix for 1. is to use proper indexing maps in the partitioner. A quick hack is to change all_users_elementwise in computation_partitioner.cc to this:

    bool all_users_elementwise =
        absl::c_all_of(instr->users(), [&](const HloInstruction* user) {
          return HloInstruction::IsOpElementwise(user->opcode()) ||
                 user->opcode() == HloOpcode::kTuple ||
                 user->opcode() == HloOpcode::kBitcast;
        });

However, I'm not sure that's safe. @pifon2a could you give that a try and see what breaks?

[pifon2a]

@jreiffers Thank you for the reproducer. Trying.

[pifon2a]

The quick hack did not work. I'll try to migrate us to the indexing maps next week.

[btaba]

Thanks @jreiffers and @pifon2a for taking a look, really appreciate it!

[kevinzakka]

+1, thank you!

[pifon2a]

So, i have a new version of a computation partitioner that completely relies on indexing maps. That did not help and it even outlines the same number of functions. Disabling the inliner helps and it compiles quickly. I will check, what's happening there.

[pifon2a]

Ok, the issue was within the inliner interface itself (XlaInlinerInterface::isLegalToInline). I will upload the fix today after checking that I did not regress anything.

[pifon2a]

The fix makes compilation of one of the tests in JAX slow. Everything else became better... Looking.

[pifon2a]

Indexing map-based partitioner and Tweaked inliner fixed the issue. Let me know if you still have problems with this.