out of memory

[slkduiers]

Traceback (most recent call last):

jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 15724724752 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
parameter allocation: 959.50MiB
constant allocation: 136.16MiB
maybe_live_out allocation: 12.84GiB
preallocated temp allocation: 14.64GiB
preallocated temp fragmentation: 480B (0.00%)
total allocation: 28.56GiB
total fragmentation: 136.17MiB (0.47%)
Peak buffers:
Buffer 1:
Size: 8.91GiB
Operator: op_name="jit()/jit(main)/while/body/remat/mesh2grid_gnn/_process/mesh2grid_gnn/_process_step/mesh2grid_gnn/concatenate[dimension=2]" source_file="/root/miniconda3/envs/large/bin/ai-models" source_line=33
XLA Label: fusion
Shape: bf16[3114720,1,1536]
==========================

    Buffer 2:
            Size: 2.97GiB
            Operator: op_name="jit(<unnamed wrapped function>)/jit(main)/while/body/remat/mesh2grid_gnn/_embed/mesh2grid_gnn/sequential/encoder_edges_mesh2grid_mlp/linear_1/dot_general[dimension_numbers=(((2,), (0,)), ((), ())) precision=None preferred_element_type=bfloat16]" source_file="/root/miniconda3/envs/large/bin/ai-models" source_line=33
            XLA Label: custom-call
            Shape: bf16[3114720,512]
            ==========================

[slkduiers]

it seem that my GPU do not have enough space to run this model

[YUTAIPAN]

Maybe you could try the larger GPU memory instance