LayerStack.from_sublayer_builder temporarily uses about twice the memory

[JEM-Mosig]

When I use the config

config = LlamalikeTransformerConfig(
    num_kv_heads=32,
    query_head_multiplier=1,
    embedding_dim=4096,
    projection_dim=128,
    mlp_hidden_dim=11008,
    num_decoder_blocks=32,
    vocab_size=32007,
    mlp_variant='swiglu',
    tie_embedder_and_logits=False,
    rope_wavelength=10000.0,
    rms_norm_eps=1e-05,
    attention_type=AttentionTypeGlobalCausal(),
    use_post_attn_norm=False,
    use_post_ffw_norm=False,
    final_logit_softcap=None,
    attn_logits_soft_cap=None,
    query_scaling_factor='default',
    parameter_dtype=jax.numpy.bfloat16,
    activation_dtype=jax.numpy.float32,
    use_layer_stack=True,
)

to run

llamalike_common.build_llamalike_transformer(config, jax.random.key(123))

I get XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 2885681152 bytes, even though the loaded model should only take up about 13 GB. When I run this on CPU, I also see the spike in memory consumption before it settles down to the appropriate value.

This does not happen when use_layer_stack=False.

[JEM-Mosig]

I managed to load the frozen layer-stacked model into GPU by stacking each parameter group on CPU and pushing it to GPU one at a time. Unfortunately, running a forward pass on a small input still results in ValueError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 1073741824 bytes.

This is on an NVIDIA L4 GPU with 24GB vRAM. After loading the model (but before the forward pass), nvidia-smi shows

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.14              Driver Version: 550.54.14      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA L4                      On  |   00000000:00:03.0 Off |                    0 |
| NA   73C    P0             35W      72W |   17058MiB /  23034MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A     26141      C   ...v0/_-q8SO6VnSpibMPUniJGc/bin/python      17050MiB |
+-----------------------------------------------------------------------------------------+

Same when I wrap the model in Jitted before the forward pass: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17180055168 bytes.

[JEM-Mosig]

Running

import os
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".98"

before any of this does not help

[JEM-Mosig]

Just passing the example input through the first layer (embedding lookup, outside the layer stack) does not result in an OOM, so it really has something to do with the LayerStack.

[danieldjohnson]

Hm, that's odd. I wonder if the initialization issue and the runtime issue have the same cause or different causes.

Do you have a traceback for when the allocation fails?

Also, does the forward pass work with use_layer_stack=False?

A possible workaround for the initialization OOM would be to wrap the initialization logic in jax.jit (or pz. variable_jit). But this won't fix the forward pass issue.

I would also be curious to see what these runs look like under profiling (https://docs.jax.dev/en/latest/profiling.html#tensorboard-profiling or https://docs.jax.dev/en/latest/device_memory_profiling.html)

[JEM-Mosig]

Thanks for the reply!

I wonder if the initialization issue and the runtime issue have the same cause or different causes.

I am unsure about this, too.

Do you have a traceback for when the allocation fails?

Here's the traceback of the forward pass without Jitted:

ValueError                                Traceback (most recent call last)
Cell In[17], line 1
----> 1 bound_model(example)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/models/transformer/model_parts.py:150, in TransformerLM.__call__(self, tokens, token_positions, **side_inputs)
    148 if token_positions is None:
    149   token_positions = pz.nx.arange("seq", tokens.named_shape["seq"])
--> 150 return self.body(tokens, token_positions=token_positions, **side_inputs)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/nn/grouping.py:77, in Sequential.__call__(self, value, **side_inputs)
     75 for i, layer in enumerate(self.sublayers):
     76   with jax.named_scope(f"{i}"):
---> 77     value = layer(value, **side_inputs)
     78 return value

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/nn/layer_stack.py:201, in LayerStack.__call__(self, argument, **side_inputs)
    193 # Separate the array leaves that have the stack axis, and move the stacked
    194 # axis to be the first positional axis (so that we can scan over it).
    195 stacked_array_selection = (
    196     selectors.select(pure_sublayer)
    197     .at_instances_of(named_axes.NamedArrayBase)
    198     .where(lambda x: self.stack_axis in x.named_shape)
    199 )
    200 named_array_slices = [
...
---> 88   outs = fun(*args)
     89 finally:
     90   lib.jax_jit.swap_thread_local_state_disable_jit(prev)

ValueError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 2885681152 bytes.

perhaps more interesting are the warning messages that appear before the error:

2025-06-07 05:41:50.760287: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.69GiB (rounded to 2885681152)requested by op 
2025-06-07 05:41:50.760440: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] ******************************************************************************************__________
E0607 05:41:50.760464   29974 pjrt_stream_executor_client.cc:2917] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 2885681152 bytes. [tf-allocator-allocation-error='']

and here with Jitted:

XlaRuntimeError                           Traceback (most recent call last)
Cell In[13], line 1
----> 1 jitted_bound_model(example)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/models/transformer/model_parts.py:150, in TransformerLM.__call__(self, tokens, token_positions, **side_inputs)
    148 if token_positions is None:
    149   token_positions = pz.nx.arange("seq", tokens.named_shape["seq"])
--> 150 return self.body(tokens, token_positions=token_positions, **side_inputs)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/toolshed/jit_wrapper.py:57, in Jitted.__call__(self, argument, **side_inputs)
     56 def __call__(self, argument: Any, /, **side_inputs) -> Any:
---> 57   return _flat_jit_call_layer(self.body, argument, side_inputs)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/core/variables.py:456, in variable_jit.<locals>.outer_fun(*args, **kwargs)
    454 (args_slots, kwargs_slots), variables = unbind_variables((args, kwargs))
    455 frozen_variables = [var.freeze() for var in variables]
--> 456 result, new_frozen_variables = jitted_inner(
    457     *args_slots, __penzai_variables=frozen_variables, **kwargs_slots
    458 )
    459 for var, new_frozen in zip(variables, new_frozen_variables):
    460   var.update(new_frozen)

    [... skipping hidden 5 frame]
...
-> 1306   results = self.xla_executable.execute_sharded(input_bufs)
   1308 if dispatch.needs_check_special():
   1309   out_arrays = results.disassemble_into_single_device_arrays()

XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17180055168 bytes.
}

preceded by

2025-06-07 05:23:12.606096: W external/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc:3022] Can't reduce memory use below 3.72GiB (3990074234 bytes) by rematerialization; only reduced to 16.52GiB (17738647836 bytes), down from 23.54GiB (25275902885 bytes) originally
2025-06-07 05:23:25.619487: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.00GiB (rounded to 17180055296)requested by op 
2025-06-07 05:23:25.619629: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] *****************************************************************************_______________________
E0607 05:23:25.619672   26141 pjrt_stream_executor_client.cc:2917] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17180055168 bytes. [tf-allocator-allocation-error='']

Also, does the forward pass work with use_layer_stack=False?

Yes, that works

I would also be curious to see what these runs look like under profiling

I'll test this soon

[danieldjohnson]

Thanks! Do you have more of the first traceback? Looks like some of the interesting parts are in the ... (in particular the part that looks like

---> 88   outs = fun(*args)
     89 finally:
     90   lib.jax_jit.swap_thread_local_state_disable_jit(prev)

seems interesting, but the part before the ... looks relevant too.)

[JEM-Mosig]

Sorry, I had not expanded the notebook cell before copying.
Using the config from my original post, running model = llamalike_common.build_llamalike_transformer(config, jax.random.key(123)) produces the warnings

2025-06-10 04:55:16.495100: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.69GiB (rounded to 2885681152)requested by op 
2025-06-10 04:55:16.495269: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] ************************_********************************************************************_______
E0610 04:55:16.495408    4857 pjrt_stream_executor_client.cc:2917] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 2885681152 bytes. [tf-allocator-allocation-error='']

followed by the traceback:

---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
Cell In[4], line 2
      1 # with jax.profiler.trace("/tmp/tensorboard"):
----> 2 model = llamalike_common.build_llamalike_transformer(config, jax.random.key(123))

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/models/transformer/variants/llamalike_common.py:475, in build_llamalike_transformer(config, init_base_rng, name)
    470   if not isinstance(config.attention_type, AttentionType):
    471     raise ValueError(
    472         "Layer stack does not currently support per-layer attention types."
    473     )
    474   sublayers.append(
--> 475       pz.nn.LayerStack.from_sublayer_builder(
    476           builder=build_llamalike_block,
    477           stack_axis="blocks",
    478           stack_axis_size=config.num_decoder_blocks,
    479           init_base_rng=init_base_rng,
    480           builder_kwargs=dict(name=f"{name}/blocks", config=config),
    481       )
    482   )
    483 else:
    484   if not isinstance(config.attention_type, AttentionType):

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/nn/layer_stack.py:322, in LayerStack.from_sublayer_builder(cls, builder, stack_axis, stack_axis_size, init_base_rng, builder_kwargs)
    316   adjusted_namedarrays = collections.OrderedDict({
    317       k: v.with_positional_prefix()
    318       for k, v in namedarray_selection.selected_by_path.items()
    319   })
    320   return namedarray_selection.remainder, adjusted_namedarrays
--> 322 remainder, namedarrays = jax.vmap(
    323     go, out_axes=(None, 0), axis_size=stack_axis_size
    324 )(
    325     None
    326     if init_base_rng is None
    327     else jax.random.split(init_base_rng, stack_axis_size)
    328 )
    329 stacked_namedarrays = {
    330     k: v.tag_prefix(stack_axis) for k, v in namedarrays.items()
    331 }
    332 stacked_sublayers, stacked_variable_values = selectors.Selection(
    333     selected_by_path=collections.OrderedDict(stacked_namedarrays),
    334     remainder=remainder,
    335 ).deselect()

    [... skipping hidden 7 frame]

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/nn/layer_stack.py:296, in LayerStack.from_sublayer_builder.<locals>.go(rng)
    295 def go(rng):
--> 296   sublayer = builder(
    297       init_base_rng=rng,
    298       **builder_kwargs,
    299   )
    300   unbound_sublayer, var_values = variables.unbind_variables(
    301       sublayer, freeze=True
    302   )
    303   if any(
    304       not isinstance(leaf, named_axes.NamedArrayBase)
    305       for leaf in jax.tree_util.tree_leaves(
    306           var_values, is_leaf=named_axes.is_namedarray
    307       )
    308   ):

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/models/transformer/variants/llamalike_common.py:413, in build_llamalike_block(name, init_base_rng, config, block_index)
    395 if config.use_post_attn_norm:
    396   attn_sequence.append(
    397       pz.nn.RMSLayerNorm.from_config(
    398           name=f"{name}/post_attention_norm",
   (...)    403       )
    404   )
    405 ffw_sequence = [
    406     pz.nn.RMSLayerNorm.from_config(
    407         name=f"{name}/pre_ffw_norm",
    408         init_base_rng=init_base_rng,
    409         across_axes={"embedding": config.embedding_dim},
    410         dtype=config.parameter_dtype,
    411         epsilon=config.rms_norm_eps,
    412     ),
--> 413     build_llamalike_feedforward(f"{name}/mlp", init_base_rng, config),
    414 ]
    415 if config.use_post_ffw_norm:
    416   ffw_sequence.append(
    417       pz.nn.RMSLayerNorm.from_config(
    418           name=f"{name}/post_ffw_norm",
   (...)    423       )
    424   )

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/models/transformer/variants/llamalike_common.py:184, in build_llamalike_feedforward(name, init_base_rng, config)
    156 else:
    157   raise ValueError(f"Unsupported MLP variant {config.mlp_variant}")
    159 return model_parts.TransformerFeedForward([
    160     pz.nn.BranchAndMultiplyTogether(
    161         branches=[
    162             pz.nn.NamedGroup(
    163                 "gate",
    164                 [
    165                     pz.nn.Linear.from_config(
    166                         name=f"{name}/gating_linear",
    167                         init_base_rng=init_base_rng,
    168                         input_axes={"embedding": config.embedding_dim},
    169                         output_axes={"neurons": config.mlp_hidden_dim},
    170                         dtype=config.parameter_dtype,
    171                     ),
    172                     pz.nn.Elementwise(act_fn),
    173                 ],
    174             ),
    175             pz.nn.Linear.from_config(
    176                 name=f"{name}/value_linear",
    177                 init_base_rng=init_base_rng,
    178                 input_axes={"embedding": config.embedding_dim},
    179                 output_axes={"neurons": config.mlp_hidden_dim},
    180                 dtype=config.parameter_dtype,
    181             ),
    182         ]
    183     ),
--> 184     pz.nn.Linear.from_config(
    185         name=f"{name}/out_linear",
    186         init_base_rng=init_base_rng,
    187         input_axes={"neurons": config.mlp_hidden_dim},
    188         output_axes={"embedding": config.embedding_dim},
    189         dtype=config.parameter_dtype,
    190     ),
    191 ])

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/nn/linear_and_affine.py:565, in Linear.from_config(cls, name, init_base_rng, input_axes, output_axes, parallel_axes, parallel_broadcast_axes, initializer, dtype, rename_outputs_if_necessary)
    552 if set(parallel_axes).intersection(
    553     set(input_axes).union(output_axes)
    554 ) or set(parallel_broadcast_axes).intersection(
    555     set(input_axes).union(output_axes, parallel_axes)
    556 ):
    557   raise ValueError(
    558       "parallel_axes and parallel_broadcast_axes must not overlap with"
    559       f" each other or with input/output axes; got input_axes={input_axes},"
    560       f" output_axes={output_axes}, parallel_axes={parallel_axes},"
    561       f" parallel_broadcast_axes={parallel_broadcast_axes}."
    562   )
    564 return cls(
--> 565     weights=parameters.make_parameter(
    566         f"{name}.weights",
    567         init_base_rng,
    568         initializer,
    569         input_axes=input_axes,
    570         output_axes=output_axes,
    571         parallel_axes={**parallel_axes, **parallel_broadcast_axes},
    572         convolution_spatial_axes={},
    573         dtype=dtype,
    574     ),
    575     in_axis_names=tuple(input_axes.keys()),
    576     out_axis_names=(
    577         tuple(output_axes.keys()) + tuple(parallel_broadcast_axes.keys())
    578     ),
    579 )

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/nn/parameters.py:112, in make_parameter(name, init_base_rng, initializer, metadata, *init_args, **init_kwargs)
    109 if metadata is None:
    110   metadata = {}
    111 return variables.Parameter(
--> 112     value=initializer(
    113         derive_param_key(init_base_rng, name), *init_args, **init_kwargs
    114     ),
    115     label=name,
    116     metadata=metadata,
    117 )

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/nn/linear_and_affine.py:152, in variance_scaling_initializer(key, scale, mode, distribution, input_axes, output_axes, parallel_axes, convolution_spatial_axes, dtype)
    146   shape.append(size)
    148 # Note: At the time of writing, JAX's variance scaling initializer requires
    149 # its input to be at least 2d, even though it's OK with passing empty tuples
    150 # for any of the axis specifications, so we have to add extra axes to the
    151 # shape to make sure this works reliably.
--> 152 array = jax.nn.initializers.variance_scaling(
    153     scale=scale,
    154     mode=mode,
    155     distribution=distribution,
    156     in_axis=input_axis_indices,
    157     out_axis=output_axis_indices,
    158     batch_axis=parallel_axis_indices,
    159 )(key=key, shape=tuple(shape) + (1, 1), dtype=dtype)[..., 0, 0]
    161 return named_axes.wrap(array).tag(*names)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:1081, in _forward_operator_to_aval.<locals>.op(self, *args)
   1080 def op(self, *args):
-> 1081   return getattr(self.aval, f"_{name}")(self, *args)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:656, in _getitem(self, item)
    655 def _getitem(self, item):
--> 656   return indexing.rewriting_take(self, item)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/numpy/indexing.py:628, in rewriting_take(arr, idx, indices_are_sorted, unique_indices, mode, fill_value, out_sharding)
    619 def rewriting_take(arr, idx, indices_are_sorted=False, unique_indices=False,
    620                    mode=None, fill_value=None, out_sharding=None):
    621   # Computes arr[idx].
   (...)    625   # For simplicity of generated primitives, we call lax.dynamic_slice in the
    626   # simplest cases: i.e. non-dynamic arrays indexed with integers and slices.
--> 628   result = _attempt_rewriting_take_via_slice(arr, idx, mode, out_sharding)
    629   if result is not None:
    630     return result

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/numpy/indexing.py:615, in _attempt_rewriting_take_via_slice(arr, idx, mode, out_sharding)
    613     arr = internal_ds(arr, start_indices)
    614 if int_indices:
--> 615   arr = lax.squeeze(arr, tuple(int_indices))
    616 return arr

    [... skipping hidden 23 frame]

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py:1306, in ExecuteReplicated.__call__(self, *args)
   1304   self._handle_token_bufs(result_token_bufs, sharded_runtime_token)
   1305 else:
-> 1306   results = self.xla_executable.execute_sharded(input_bufs)
   1308 if dispatch.needs_check_special():
   1309   out_arrays = results.disassemble_into_single_device_arrays()

XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 2885681152 bytes.

The traceback you asked about happens when I circumvent the model loading OOM with the stack-on-CPU-trick and run model(example):

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[17], line 1
----> 1 bound_model(example)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/models/transformer/model_parts.py:150, in TransformerLM.__call__(self, tokens, token_positions, **side_inputs)
    148 if token_positions is None:
    149   token_positions = pz.nx.arange("seq", tokens.named_shape["seq"])
--> 150 return self.body(tokens, token_positions=token_positions, **side_inputs)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/nn/grouping.py:77, in Sequential.__call__(self, value, **side_inputs)
     75 for i, layer in enumerate(self.sublayers):
     76   with jax.named_scope(f"{i}"):
---> 77     value = layer(value, **side_inputs)
     78 return value

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/nn/layer_stack.py:201, in LayerStack.__call__(self, argument, **side_inputs)
    193 # Separate the array leaves that have the stack axis, and move the stacked
    194 # axis to be the first positional axis (so that we can scan over it).
    195 stacked_array_selection = (
    196     selectors.select(pure_sublayer)
    197     .at_instances_of(named_axes.NamedArrayBase)
    198     .where(lambda x: self.stack_axis in x.named_shape)
    199 )
    200 named_array_slices = [
--> 201     named_array.untag_prefix(self.stack_axis).with_positional_prefix()
    202     for named_array in stacked_array_selection.get_sequence()
    203 ]
    205 def body(carry_data, slice_data):
    206   cur_arg, shared_var_states = carry_data

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/penzai/core/named_axes.py:1695, in NamedArrayView.with_positional_prefix(self)
   1692   return NamedArray(data_array=self.data_array, named_axes=named_axes)
   1693 else:
   1694   return NamedArray(
-> 1695       data_array=self.data_array.transpose(transposition),
   1696       named_axes=named_axes,
   1697   )

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:427, in _transpose(self, *args)
    425 else:
    426   axis = _ensure_index_tuple(args)
--> 427 return lax_numpy.transpose(self, axis)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:1211, in transpose(a, axes)
   1209 axes_ = list(range(a.ndim)[::-1]) if axes is None else axes
   1210 axes_ = [_canonicalize_axis(i, np.ndim(a)) for i in axes_]
-> 1211 return lax.transpose(a, axes_)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/lax/lax.py:2915, in transpose(operand, permutation)
   2912   return operand
   2913 else:
-> 2915   return transpose_p.bind(operand, permutation=permutation)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/core.py:531, in Primitive.bind(self, *args, **params)
    529 def bind(self, *args, **params):
    530   args = args if self.skip_canonicalization else map(canonicalize_value, args)
--> 531   return self._true_bind(*args, **params)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/core.py:551, in Primitive._true_bind(self, *args, **params)
    549 trace_ctx.set_trace(eval_trace)
    550 try:
--> 551   return self.bind_with_trace(prev_trace, args, params)
    552 finally:
    553   trace_ctx.set_trace(prev_trace)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/core.py:556, in Primitive.bind_with_trace(self, trace, args, params)
    555 def bind_with_trace(self, trace, args, params):
--> 556   return trace.process_primitive(self, args, params)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/core.py:1060, in EvalTrace.process_primitive(self, primitive, args, params)
   1058 args = map(full_lower, args)
   1059 check_eval_args(args)
-> 1060 return primitive.impl(*args, **params)

File ~/2025Q2_v4/mlres/.venv/lib/python3.12/site-packages/jax/_src/dispatch.py:88, in apply_primitive(prim, *args, **params)
     86 prev = lib.jax_jit.swap_thread_local_state_disable_jit(False)
     87 try:
---> 88   outs = fun(*args)
     89 finally:
     90   lib.jax_jit.swap_thread_local_state_disable_jit(prev)

ValueError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 2885681152 bytes.

[JEM-Mosig]

I'm not very well versed in profiling GPUs, but I'm happy to learn this stuff now! I installed tensorboard and ran this:

os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".98"

config = llamalike_common.LlamalikeTransformerConfig(
    num_kv_heads=32,
    query_head_multiplier=1,
    embedding_dim=4096,
    projection_dim=128,
    mlp_hidden_dim=11008,
    # num_decoder_blocks=32,   # <-- I don't get useful infos in the profiler when OOM happens
    num_decoder_blocks=10,     # <-- so I reduce the number of blocks to avoid OOM
    vocab_size=32007,
    mlp_variant='swiglu',
    tie_embedder_and_logits=False,
    rope_wavelength=10000.0,
    rms_norm_eps=1e-05,
    attention_type=llamalike_common.AttentionTypeGlobalCausal(),
    use_post_attn_norm=False,
    use_post_ffw_norm=False,
    final_logit_softcap=None,
    attn_logits_soft_cap=None,
    query_scaling_factor='default',
    parameter_dtype=jax.numpy.bfloat16,
    activation_dtype=jax.numpy.float32,
    use_layer_stack=True,
)

# This no longer leads to OOM - did the tensorboard installation change something?
model = llamalike_common.build_llamalike_transformer(config, jax.random.key(123))

with jax.profiler.trace("/tmp/tensorboard"):
    # This leads to OOM when num_decoder_blocks=32, but I only use 10 here
    output = model(pz.nx.wrap(jnp.asarray([1, 2, 3]), "seq"))
    output.data_array.block_until_ready()

The last block prints

2025-06-13 09:51:22.401766: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1749808282.425879    8713 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1749808282.433496    8713 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
W0000 00:00:1749808282.452128    8713 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
W0000 00:00:1749808282.452148    8713 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
W0000 00:00:1749808282.452150    8713 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
W0000 00:00:1749808282.452152    8713 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.

as well as five copies of

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
AttributeError: 'MessageFactory' object has no attribute 'GetPrototype'

This graph looks suspicious to me:

[JEM-Mosig]

Here's the full profile:

profile.tar.gz

This was with a single L4 GPU and
'jax': '0.6.1',
'jaxlib': '0.6.1',
'penzai': '0.2.5',

[JEM-Mosig]

This is the memory profile with the same code, but use_layer_stack=False:

[JEM-Mosig]

When I downgrade to jax 0.5.0 (or 0.4.38) and enable use_layer_stack=True again, the graph has a much sharper spike:

When I use 32 layers, this still results in an OOM.

[JEM-Mosig]

jax 0.5.3 still gives the slow-rising triangle memory graph like 0.6.1

[danieldjohnson]

Interesting, thanks for digging into it!

From the stacktraces you shared in #122 (comment), my best guess about what's going on is:

• During initialization, the large array containing all of the parameters is initialized with two dimensions of size 1, and then these dimensions are sliced away. Ideally this shouldn't use double the memory, but it's possible JAX isn't able to share the memory of these for some reason and instead returns a copy.
  • You can probably work around this by JIT-compiling the initializer. It might also work to jit-compile penzai.nn.linear_and_affine.variance_scaling_initializer itself, which should fix it even if you initialize outside of JIT at the top level.
• For the call, it looks like the OOM is happening when we try to transpose the stack axis to the front in order to scan over it. Do you happen to know in your case whether the stack axis is already at the beginning in the data array? (E.g. if you look at the fields of the NamedArray, does the stack axis name appear first?)
  • If it is already at the front, then maybe there's an unnecessary transpose happening here. I had hoped this would again not use any extra memory, but maybe transposing an array with the identity transposition sometimes results in a copy?
  • If it's not already at the front, maybe we could fix this by moving it to the front when the layer is constructed.

[JEM-Mosig]

I'm a bit confused about the OOM during initialization happening only sometimes. It must depend on something in the environment, but I haven't figured out yet what it is. I can always circumvent it with the build-on-cpu-then-push-to-gpu-trick, though. So the forward pass issue is more severe.

I'll have a look at the axes order again later.