Explanation BufferAssignment OOM Debugging

[JohannesEsslinger]

Hi!

My question relates to OOM memory debugging in Jax. Unfortunately there is not much information in the documentation or other sources. Unfortunately similar thread, which can be found here has not yet been addressed.

I have the problem that i get the following OOM expection, but don't know why 15! buffers (for the same function call) are allocated.

jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 792000000 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:  323.89MiB
              constant allocation:        53B
        maybe_live_out allocation:    6.26GiB
     preallocated temp allocation:  521.91MiB
                 total allocation:    7.09GiB
Peak buffers:
	Buffer 1:
		Size: 755.31MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200 deduplicated_name="fusion.42"
		XLA Label: fusion
		Shape: f32[11,1,1,9,2000,1000]
		==========================

	Buffer 2:
		Size: 755.31MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200 deduplicated_name="fusion.42"
		XLA Label: fusion
		Shape: f32[11,1,1,9,2000,1000]
		==========================

	Buffer 3:
		Size: 755.31MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200 deduplicated_name="fusion.42"
		XLA Label: fusion
		Shape: f32[11,1,1,9,2000,1000]
		==========================

	Buffer 4:
		Size: 755.31MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200 deduplicated_name="fusion.42"
		XLA Label: fusion
		Shape: f32[11,1,1,9,2000,1000]
		==========================

	Buffer 5:
		Size: 755.31MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200 deduplicated_name="fusion.42"
		XLA Label: fusion
		Shape: f32[11,1,1,9,2000,1000]
		==========================

	Buffer 6:
		Size: 755.31MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200
		XLA Label: fusion
		Shape: f32[11,1,1,9,2000,1000]
		==========================

	Buffer 7:
		Size: 83.92MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200
		XLA Label: fusion
		Shape: f32[11,2000,1000]
		==========================

	Buffer 8:
		Size: 83.92MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200
		XLA Label: fusion
		Shape: f32[11,2000,1000]
		==========================

	Buffer 9:
		Size: 83.92MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200
		XLA Label: fusion
		Shape: f32[11,2000,1000]
		==========================

	Buffer 10:
		Size: 83.92MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200 deduplicated_name="fusion.49"
		XLA Label: fusion
		Shape: f32[11,1,1,2000,1000]
		==========================

	Buffer 11:
		Size: 83.92MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200
		XLA Label: fusion
		Shape: f32[11,1,1,2000,1000]
		==========================

	Buffer 12:
		Size: 83.92MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200
		XLA Label: fusion
		Shape: f32[11,1,1,2000,1000]
		==========================

	Buffer 13:
		Size: 83.92MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200 deduplicated_name="fusion.49"
		XLA Label: fusion
		Shape: f32[11,1,1,2000,1000]
		==========================

	Buffer 14:
		Size: 83.92MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200 deduplicated_name="fusion.49"
		XLA Label: fusion
		Shape: f32[11,1,1,2000,1000]
		==========================

	Buffer 15:
		Size: 83.92MiB
		Operator: op_name="jit(scan)/jit(main)/while/body/dynamic_update_slice" source_file="/XXX.py" source_line=200
		XLA Label: fusion
		Shape: f32[11,1,2000,1,1000]
		==========================

The function in source_line 200 is:

`unused_final_state, per_step_outputs = jax.lax.scan(scan_step, initial_state, per_step_inputs)`

whereby per_step_input is a tuple (dict, ArrayImpl:(11, 2), ArrayImpl(11,)).

I have to mention that 11 is the number of iterations, 2000 batch size, and 1000 the dimension of a single sample. (Dunno where the 9 is coming from)

Do you have any guess what could cause Jax to create 15 buffer or have any suggestions to debug OOM?

Kind regards!