How to deal with excessive fragmentation on TPU?

[dlwh]

Hi,

I have a computation (basically just pjit gpt-2) that by rights should fit in 16GB of HBM, but it doesn't due to fragmentation in excess of 50%.

jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Ran out of memory in memory space hbm. Used 17.97G of 15.48G hbm. Exceeded hbm capacity by 2.49G.

Total hbm usage >= 18.49G:
    reserved        530.00M
    program           9.20G
    arguments         8.77G

Output size 8.77G; shares 8.77G with arguments.

Program hbm requirement 9.20G:
    global           436.0K
    scoped            1.64M
    HLO temp          9.19G (98.7% utilization: Unpadded (4.43G) Padded (4.49G), 51.2% fragmentation (4.70G))

If i could get ~half of the fragmentation to go away, I'd be good. There are a few sizeable unnecessary buffers being allocated, and I'll try to make those disappear, but those add up to ~1.4G. (So those might be enough if it helps fragmentation...)

[skye]

We're working on automatic defragmentation, but in the meantime, try this secret menu item:
jax.lib.xla_bridge.get_backend().defragment()

This will automatically wait for all jax operations to finish, but isn't thread-safe, so don't call it if you have jax operations running in different threads. Hope this helps!

[skye]

@hawkinsp points out it's saying that the program + arguments don't fit (9.20G + 8.77G), and that the program is really large due to "HLO temp". I believe the fragmentation is somehow inside that HLO temp space itself.

I need to check with some TPU experts to better interpret this, but in the meantime, any chance you could provide us with:

1. The HLO that produced this error
2. /tmp/tpu_logs/tpu_driver.INFO immediately after getting the error

To get the HLO, set the env var: XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH". It'll produce a bunch of HLO files, you can just zip up all of them if that's easiest.

[dlwh]

Thanks so much! Here it is. gpt2_xla_dumps.tar.gz The op is module_0163.pjit_train_step

Possibly related and possibly not (and I could be totally off base!), but the one piece of (seemingly) unnecessary memory hunger that I alluded to in my first message is maybe due to the system pre-converting ~all of my parameters (stored in f32) to bf16 in one block at the beginning (lines 4537-4548 of xla_dumps/module_0163.pjit_train_step.13.after_optimizations_after_buffer_assignment.txt). I do want to do the computations with these in bf16, but doing it all up front uses a lot of memory[*]. I haven't figured out why it's doing the conversion so soon or how to coax it into not doing it in bulk.

[*] at least ~1.4G nominally, though it might all get lost to fragmentation or alternatively could save more thanks to fragmentation?

[skye]

Thanks! I'll fwd this info to the TPU team.

They've advised setting this env var: LIBTPU_INIT_ARGS=--xla_hlo_scheduling_brkga_enable_as_fallback=true
Let us know if that makes any difference?

[dlwh]

I get ERROR: Unknown command line flag 'xla_hlo_scheduling_brkga_enable_as_fallback'

When I run strings on libtpu.so I see:

$ strings /files/venv310/libtpu.so/libtpu.so  | grep brkga
xla_hlo_scheduling_brkga_compute_runtime_estimates
brkga
hlo-memory-scheduler-with-brkga-fallback
brkga_generation_limit
brkga

Based on that, I tried using hlo-memory-scheduler-with-brkga-fallback, --hlo_memory_scheduler_with_brkga_fallback, 'xla_hlo_memory_scheduler_with_brkga_fallback', and xla-hlo-memory-scheduler-with-brkga-fallback and none of those worked either.

[rjpower]

Old issue, but the BRKGA scheduler was disabled when we initially rolled out libtpu.

At the time it was considered low-importance and somewhat hard to refactor: I'm not sure how often it's useful in practice...