Why is my JAX profiler not displaying the full function stack?

[MelodicDrumstep]

Hello,

I’ve recently been working on profiling and tuning the inference procedure for the AlphaFold3 model. As someone relatively new to JAX, I’ve been experimenting with the JAX profiler as outlined in the official documentation:

  with profiler.trace(my_path_to_the_trace_result, create_perfetto_link=True):
    // SOME WORK HERE

However, the results I’ve obtained have been somewhat disappointing:

The host running time is predominantly dominated by "pjit helper," which obscures the actual function calls. Additionally, the device running time is not fully displayed. I’ve reviewed the troubleshooting section in the official documentation, but my issue doesn’t seem to align with the cases mentioned there.

For context, the process doesn’t output any error messages related to this issue. The full output is as follows:

2025-01-25 16:04:44.463272: W external/xla/xla/service/gpu/nvptx_compiler.cc:930] The NVIDIA driver's CUDA version is 12.3 which is older than the PTX compiler version 12.6.77. Because the driver is older than the PTX compiler version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.
Skipping data pipeline...
Output directory: /dssg/home/acct-hpc/alphafold3/output_dir/37aa_2jo9
Writing model input JSON to /dssg/home/acct-hpc/alphafold3/output_dir/37aa_2jo9
Predicting 3D structure for 37aa_2JO9 for seed(s) (1,)...
Featurising data for seeds (1,)...
Featurising 37aa_2JO9 with rng_seed 1.
I0125 16:04:51.946228 22894668551168 pipeline.py:160] processing 37aa_2JO9, random_seed=1
Featurising 37aa_2JO9 with rng_seed 1 took 0.58 seconds.
Featurising data for seeds (1,) took  4.92 seconds.
2025-01-25 16:04:52.848254: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1737792292.862857 1839003 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1737792292.867310 1839003 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
Running model inference for seed 1...
Running model inference for seed 1 took  56.73 seconds.
Extracting output structures (one per sample) for seed 1...
/dssg/home/acct-hpc/hpc/apps/miniconda3/envs/af3_exp/lib/python3.11/site-packages/alphafold3/model/confidences.py:332: RuntimeWarning: invalid value encountered in divide
  return np.nanmean(value * mask_with_nan, axis=axis) / np.nanmean(
/dssg/home/acct-hpc/hpc/apps/miniconda3/envs/af3_exp/lib/python3.11/site-packages/alphafold3/model/confidences.py:522: RuntimeWarning: Mean of empty slice
  xchain = np.nanmean(
/dssg/home/acct-hpc/hpc/apps/miniconda3/envs/af3_exp/lib/python3.11/site-packages/alphafold3/model/confidences.py:548: RuntimeWarning: Mean of empty slice
  return np.nanmean(np.stack([xchain_row_agg, xchain_col_agg], axis=0), axis=0)
Extracting output structures (one per sample) for seed 1 took  0.16 seconds.
Running model inference and extracting output structures for seed 1 took  56.89 seconds.
I0125 16:06:20.198991 22894668551168 profiler.py:145] Loading trace.json.gz and removing its metadata...
I0125 16:06:21.838346 22894668551168 profiler.py:154] Writing perfetto_trace.json.gz...
Open URL in browser: https://ui.perfetto.dev/#!/?url=http://127.0.0.1:9001/perfetto_trace.json.gz

In an attempt to better understand the issue, I conducted an experiment where I removed a JIT statement from a core function:

    # Before
    #return functools.partial(
    #    jax.jit(forward_fn.apply, device=self._device), self.model_params
    #)
    # After
    return functools.partial(
        forward_fn.apply, self.model_params
    )

The results after this change were as follows:

The host and device profiles now display more information and are more readable. However, I’m puzzled by the host profile output, which suggests that model.py:_call__ doesn’t call further functions for an extended period (which contradicts the actual code). Additionally, in the device profile output, the kernel function elapsed time appears unusually short (though this might be normal).

Thank you in advance! Here are my profile output files:
perfetto_trace in the first case
perfetto_trace in the second case