Correct way to perform line profiling

[shyams2]

I'm trying to analyze some code using line_profiler to understand which functions are the most expensive / identify possible speed-ups. However, I'm uncertain how we profiling should be performed when we're also jit-compiling. For instance, if I wrap the function as:

@profile
@jax.jit
def step(...):
    # your code here

it throws an error AttributeError: 'CompiledFunction' object has no attribute '__code__'

However, instead, if I use

@jax.jit
@profile
def step(...):
    # your code here

it seems to run. However, I'm not sure if it's only profiling the compilation time.
Is there a better way by which I can profile to understand bottlenecks better?

Output Observed

To provide more details about the specific output I'm seeing : Here's the output of the line profiler I'm seeing in the training loop when I'm calling the function I want to profile:

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
 373     10002  495914672.0  49581.6     97.8              self.params, self.opt_state = self.step(
   374      5001       2417.0      0.5      0.0                  self.params, self.opt_state, batch
   375                                                       )

Here it seems like the time spent in the step function is 495s. And the corresponding output for the step function:

Total time: 21.4522 s
File: benchmark_singlenode_singleGPU.py
Function: step at line 312

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
   312                                               @partial(jax.jit, static_argnums=(0,))
   313                                               @profile
   314                                               def step(self, params, opt_state, batch):
   315         1   21307501.0 21307501.0     99.3          g = grad(self.loss)(params, batch)
   316         1     134959.0 134959.0      0.6          updates, opt_state = self.optimizer.update(g, opt_state)
   317         1       9756.0   9756.0      0.0          params = optax.apply_updates(params, updates)
   318         1          1.0      1.0      0.0          return params, opt_state

Here, you can see that it states that total time spent is 21s which seems to suggest that this profiling is just for the compilation stage.

[jakevdp]

Under jit, the code withing a function is being compiled and run as a unit rather than being run by the python interpreter, so typical Python line-profiling approaches will not work. You'll have to use other tools: the documentation has some information on profiling here: https://jax.readthedocs.io/en/latest/profiling.html

[hawkinsp]

fusions are kernels generated by XLA. cutlass is one of Nvidia's libraries for matrix multiplications.

As to why the device trace is so short, try only sampling a short period (e.g., 10s or so). I think it's possible one of the device profiling buffers is overflowing and dropping GPU events because your trace is very long.

[shyams2]

@hawkinsp I see. Thanks a lot for the clarification. It definitely seems like sampling for a shorter period works a lot better. I've tried doing the profiling on nsys instead. However, there's one thing that isn't fully clear to me yet. I'm seeing that a lot of time is spent on launching kernels. Is this normal behaviour for deep learning codes? Any advice on how one would go about reducing these kernel launches in general? (data from nsys)

[hawkinsp]

Personally, I'd look at a trace view rather than a profile view to determine whether host-side events (e.g., CUDA kernel launches) are actually worth optimizing. If your GPU is being kept busy all of the time, it doesn't matter greatly what the host overheads are most likely.

[coreyjadams]

I'm curious, based on this discussion: when you JIT a function, does that fuse into one GPU kernel or is it a sequence of kernels in XLA? I understand your point on the trace viewing, but in the event that someone wanted to reduce the kernel launches, does bigger JIT'd functions do the trick?

[hawkinsp]

On GPU, it ends up being a sequence of kernels. Exactly how your operations map to kernels is up to XLA.

There are a number of reasons you don't just end up with one gigantic kernel. One is that we want to call NVidia-provided kernels, e.g., matmul. Another is that different computations are going to have different thread block structures, so they need to go in separate kernels.