XLA duplicates input arguments memory because of wrapped_transpose

[qGentry]

Hi, I have simple LLM training step that computes forward+backward pass and updates weights. It is LoRA training, so only very small subset of parameters is trainable (70B total parameters, ~600M trainable). I've checked memory analysis and found that actual memory consumption is way bigger than expected and figured out that this is because XLA is doing some very weird transposition of input weights, resulting in effectively doubling of memory usage for 70B model parameters.

Can I somehow disabled this optimization or overcome this issue in some other way?

I've tested it on both JAX 0.4.34 and latest version (JAX 0.6.1)

Here's attached HLO dump:

hlo.tgz

Some of the transposed buffer names:

wrapped_transpose.5
input_transpose_fusion.3
wrapped_transpose.10

[qGentry]

This may be somewhat related to
jax-ml/jax#20303

[mooskagh]

I cannot come up with a simple way to fix it (rather than pre-transpose weights), but here's what happens:

Some ops, like dynamic slices, require tensor to be transposed a certain way, for example this one:

  ROOT %dynamic-slice.318.1 = bf16[1024,1,29568]{2,1,0} dynamic-slice(%bitcast.780.1, %constant_1683_24, %param_0.1196, %constant_1683_24), dynamic_slice_sizes={1024,1,29568}, metadata={op_name="jit(<unnamed wrapped function>)/jit(main)/while/body/jvp(CausalLMTransformerModel.compute_logits)/transformer/while/body/dynamic_slice" source_file="/usr/local/lib/python3.10/dist-packages/flax/core/axes_scan.py" source_line=179}

It's located inside a loop (actually, two nested loops, right?), so optimizer decides to hoist that transposition out of the loop. As a result, the transposed copy of weights have to be alive through the entire duration of the loop. And that's for all weights.

One thing that sometimes helps is AUTO layout (e.g. https://github.com/jax-ml/jax/blob/main/tests/layout_test.py#L36) which would do transposes on the host, but I'm not sure how easy it is to make it work.

[qGentry]

Yes, this is indeed a double while loop (gradient accumulation and scan over layers).
We have quite a complex sharding structure so I don't feel like auto layout might work for me.

[qGentry]

Here's also an HLO with --xla_dump_hlo_pass_re=.*

hlo_with_re.tgz

[mooskagh]

In the "with_re" archive, I don't see any wrapped transposes, is it the same model (also the file itself is smaller)?

Could you try --xla_disable_hlo_passes=while-loop-invariant-code-motion?
I'm not sure that that's the pass that does it though, I can check later (probably tomorrow) which exactly one is it.

Other options to try meanwhile are while-loop-constant-sinking and simplify-while-loops (which from reading don't seem like they do that, but other passes don't seem to be active in GPU pipeline)

[qGentry]

Yeah, the issue with previous HLO was in compilation cache. Here's one with it. Couldn't attach it to github because of file size restrictions.
https://drive.google.com/file/d/1H_BCJ7vojqWQM2wbOVlLOB_k-NQUxj20/view?usp=sharing

[qGentry]

I've tried every flag but neither of them had any effect on memory usage

2025-05-26 15:04:40 INFO:root:XLA flags set: --xla_disable_hlo_passes=rematerialization,while-loop-invariant-code-motion,while-loop-constant-sinking,simplify-while-loops --xla_gpu_graph_level=0 --xla_gpu_enable_triton_gemm=false

[jreiffers]

This appears to be caused by layout assignment, which introduces a transposing copy in the top level. Not sure yet why that's happening, we're having some trouble with VLOGing.

[mooskagh]

Ah yeah, that makes sense, we've seen that quite a few times.

Layout assignment goes from inner computations to outer, and then in the entry computation when there's a clash between parameter layout and nested computation parameter layout, there's a transpose.

There's no easy way to fix it correctly, as all computations are processed in isolation (and sometimes it's the model parameter that can "adsorb" the transpose by being auto layout, sometimes it's an insturction in the nested computation that doesn't care), but I'll think about workarounds.

For the OOM issue in particular, I believe we should have a pass quite late in the pipeline that would sink operations into the loop body to reduce the lifetime of buffers (certainly that doesn't go as a quick workaround).

[jreiffers]

AFAICT, the reason are all these all-gathers with dimension=1, like this:

%all-gather.41 = bf16[1,64,8192]{2,1,0} all-gather(%convert.442), channel_id=61, replica_groups=[1,8]<=[8], dimensions={1}, use_global_device_ids=true

Because the all-gather dimension has to be major, gpu_layout_assignment thinks that they need to get a layout of {2,0,1}, which then gets propagated everywhere. But I don't think that's actually true, because the shape is [1,64,8192], i.e. dim 0 is degenerate, so {2,1,0} should work as well and everything should be fine.

@mooskagh I see a couple of ways to fix this, e.g. by getting rid of these degenerate dims or by attempting to make gpu_layout_assignment smarter. The former seems easier to me, WDYT?

[jreiffers]

4c817c5 This gets rid of most of the copies, a handful are still there. Looking into the remaining ones now.

[jreiffers]

The remaining copies are actually needed, because the all-gathers are gathering along a minor dimension: