XLA attempts to broadcast weights to batch size

[neel04]

jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=85899345920) would exceed memory (size=34359738368) :: #allocation10824 [shape = 'f32[512,40,1024,1024]{3,2,0,1:T(8,128)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of
broadcast.65605.clone.5@{}'] :: <no-hlo-instruction>

Effectively I'm doing DDP and have a setup where the {query | key | value} projection is of shape (40, 1024, 1024) and the batch size of my data is 512. So it seems that for attention projections, its broadcasting the the weights over the data dimension for some reason, which leads to the OOM.

Strangely enough, that's not happening on other layers - such as embedding which if broadcasted would also lead to an OOM.

I can verify I'm sharding my data on the 0th axis, and replicating my model correctly so it's strange.

    ...
    def get_devices(self, model_axis = 1):
        return mesh_utils.create_device_mesh(
            (jax.device_count() // model_axis, model_axis),
            allow_split_physical_axes=True,
        )

    def get_mesh(self) -> Mesh:
        return Mesh(self.get_devices(), axis_names=('data', 'model'))

    def shard_data(self, tree: PyTree | Array) -> PyTree | Array:
        return eqx.filter_shard(tree, NamedSharding(self.mesh, P('data')))

    def shard_model(self, tree: PyTree) -> PyTree:
        return jtu.tree_map(self.ddp_sharding, tree)

    def shard_one_hot(self, tree: PyTree) -> PyTree:
        return tree
        
    def ddp_sharding(self, leaf: PyTree) -> PyTree:
        if not eqx.is_array(leaf):
            return leaf

        sharding_ = NamedSharding(self.mesh, P())

        return eqx.filter_shard(leaf, sharding_)

Any pointers how I can approach debugging? I've been poking around sharding but it seems alright and my attention is definitely under vmap, so any broadcasting happening is done by XLA behind the scenes.

What could be the possible points of failure I should check?

Full traceback

Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/optuna/study/_optimize.py", line 197, in _run_trial
    value_or_values = func(trial)
                      ^^^^^^^^^^^
  File "/ReAct_Jax/train_model.py", line 118, in <lambda>
    lambda trial: kickoff_optuna(trial=trial, **trainer_kwargs),
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/ReAct_Jax/train_model.py", line 201, in kickoff_optuna
    output = trainer.train(trial)
             ^^^^^^^^^^^^^^^^^^^^
  File "/ReAct_Jax/ReAct/utils/trainer.py", line 381, in train
    loss, model, opt_state, grads, updates = make_step(
                                             ^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/equinox/_jit.py", line 242, in __call__
    return self._call(False, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/equinox/_module.py", line 1078, in __call__
    return self.__func__(self.__self__, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/equinox/_jit.py", line 213, in _call
    out = self._cached(dynamic_donate, dynamic_nodonate, static)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 332, in cache_miss
    outs, out_flat, out_tree, args_flat, jaxpr, attrs_tracked = _python_pjit_helper(
                                                                ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 190, in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **p.params)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/core.py", line 2739, in bind
    return self.bind_with_trace(top_trace, args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/core.py", line 433, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/core.py", line 939, in process_primitive
    return primitive.impl(*tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 1730, in _pjit_call_impl
    return xc._xla.pjit(
           ^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 1712, in call_impl_cache_miss
    out_flat, compiled = _pjit_call_impl_python(
                         ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 1642, in _pjit_call_impl_python
    ).compile(compile_options)
      ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2295, in compile
    executable = UnloadedMeshExecutable.from_hlo(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2807, in from_hlo
    xla_executable = _cached_compilation(
                     ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2621, in _cached_compilation
    xla_executable = compiler.compile_or_get_cached(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/compiler.py", line 399, in compile_or_get_cached
    return _compile_and_write_cache(
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/compiler.py", line 627, in _compile_and_write_cache
    executable = backend_compile(
                 ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/profiler.py", line 336, in wrapper
   File "/usr/local/lib/python3.11/site-packages/jax/_src/profiler.py", line 336, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/compiler.py", line 267, in backend_compile
    return backend.compile(built_c, compile_options=options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=85899345920) would exceed memory (size=34359738368) :: #allocation10824 [shape = 'f32[512,40,1024,1024]{3,2,0,1:T(8,128)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of
broadcast.65605.clone.5@{}'] :: <no-hlo-instruction>
[W 2024-09-07 00:13:29,003] Trial 0 failed with value None.
Traceback (most recent call last):
  File "/ReAct_Jax/train_model.py", line 207, in <module>
    main(key)
  File "/ReAct_Jax/train_model.py", line 117, in main
    study.optimize(
  File "/usr/local/lib/python3.11/site-packages/optuna/study/study.py", line 475, in optimize
    _optimize(
  File "/usr/local/lib/python3.11/site-packages/optuna/study/_optimize.py", line 63, in _optimize
    _optimize_sequential(
  File "/usr/local/lib/python3.11/site-packages/optuna/study/_optimize.py", line 160, in _optimize_sequential
    frozen_trial = _run_trial(study, func, catch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/optuna/study/_optimize.py", line 248, in _run_trial
    raise func_err
  File "/usr/local/lib/python3.11/site-packages/optuna/study/_optimize.py", line 197, in _run_trial
    value_or_values = func(trial)
                      ^^^^^^^^^^^
  File "/ReAct_Jax/train_model.py", line 118, in <lambda>
    lambda trial: kickoff_optuna(trial=trial, **trainer_kwargs),
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/ReAct_Jax/train_model.py", line 201, in kickoff_optuna
    output = trainer.train(trial)
             ^^^^^^^^^^^^^^^^^^^^
  File "/ReAct_Jax/ReAct/utils/trainer.py", line 381, in train
    loss, model, opt_state, grads, updates = make_step(
                                             ^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/equinox/_jit.py", line 242, in __call__
    return self._call(False, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/equinox/_module.py", line 1078, in __call__
    return self.__func__(self.__self__, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/equinox/_jit.py", line 213, in _call
    out = self._cached(dynamic_donate, dynamic_nodonate, static)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 332, in cache_miss
    outs, out_flat, out_tree, args_flat, jaxpr, attrs_tracked = _python_pjit_helper(
                                                                ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 190, in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **p.params)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/core.py", line 2739, in bind
    return self.bind_with_trace(top_trace, args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/core.py", line 433, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/core.py", line 939, in process_primitive
    return primitive.impl(*tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 1730, in _pjit_call_impl
    return xc._xla.pjit(
           ^^^^^^^^^^^^^  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 1712, in call_impl_cache_miss
    out_flat, compiled = _pjit_call_impl_python(
                         ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/pjit.py", line 1642, in _pjit_call_impl_python
    ).compile(compile_options)
      ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2295, in compile
    executable = UnloadedMeshExecutable.from_hlo(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2807, in from_hlo
    xla_executable = _cached_compilation(
                     ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2621, in _cached_compilation
    xla_executable = compiler.compile_or_get_cached(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/compiler.py", line 399, in compile_or_get_cached
    return _compile_and_write_cache(
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/compiler.py", line 627, in _compile_and_write_cache
    executable = backend_compile(
                 ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/profiler.py", line 336, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/jax/_src/compiler.py", line 267, in backend_compile
    return backend.compile(built_c, compile_options=options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=85899345920) would exceed memory (size=34359738368) :: #allocation10824 [shape = 'f32[512,40,1024,1024]{3,2,0,1:T(8,128)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of
broadcast.65605.clone.5@{}'] :: <no-hlo-instruction>

Model `PyTree`

Model # of parameters: 609.03M
Model: React(
  max_iters=3,
  width=1024,
  embed_layer=Embedding(
    num_embeddings=50304,
    embedding_size=1024,
    weight=f32[50304,1024]
  ),
  embed_ln=LayerNorm(
    shape=(1024,),
    eps=1e-05,
    use_weight=True,
    use_bias=True,
    weight=f32[1024],
    bias=f32[1024]
  ),
  main_block=RecurrentModule(
    num_layers=40,
    attention_layers=AttentionBlock(
      seqlen=512,
      n_heads=8,
      in_dim=1024,
      attn_gate=MultiheadAttention(
        query_proj=Linear(
          weight=f32[40,1024,1024],
          bias=f32[40,1024],
          in_features=1024,
          out_features=1024,
          use_bias=True
        ),
        key_proj=Linear(
          weight=f32[40,1024,1024],
          bias=f32[40,1024],
          in_features=1024,
          out_features=1024,
          use_bias=True
        ),
        value_proj=Linear(
          weight=f32[40,1024,1024],
          bias=f32[40,1024],
          in_features=1024,
          out_features=1024,
          use_bias=True
        ),
        output_proj=Linear(
          weight=f32[40,1024,1024],
          bias=f32[40,1024],
          in_features=1024,
          out_features=1024,
          use_bias=True
        ),
        dropout=Dropout(p=0.01, inference=False),
        num_heads=8,
        query_size=1024,
        key_size=1024,
        value_size=1024,
        output_size=1024,
        qk_size=128,
        vo_size=128,
        use_query_bias=True,
        use_key_bias=True,
        use_value_bias=True,
        use_output_bias=True
      ),
      rope_embed=RotaryPositionalEmbedding(embedding_size=128, theta=10000.0),
      ln1=LayerNorm(
        shape=(1024,),
        eps=1e-05,
        use_weight=True,
        use_bias=True,
        weight=f32[40,1024],
        bias=f32[40,1024]
      ),
      ln2=LayerNorm(
        shape=(1024,),
        eps=1e-05,
        use_weight=True,
        use_bias=True,
        weight=f32[40,1024],
        bias=f32[40,1024]
      ),
      mlp=MLP(
        layer_1=LinearProj(
          bias=f32[40,4096],
          weight=f32[40,1024,4096],
          input_dim=1024,
          output_dim=4096,
          use_bias=True
        ),
        layer_2=LinearProj(
          bias=f32[40,1024],
          weight=f32[40,4096,1024],
          input_dim=4096,
          output_dim=1024,
          use_bias=True
        ),
        dropout=Dropout(p=0.01, inference=False),
        act=NewGELU()
      )
    ),
    post_ln=LayerNorm(
      shape=(1024,),
      eps=1e-05,
      use_weight=True,
      use_bias=True,
      weight=f32[1024],
      bias=f32[1024]
    ),
    reshape_gate=LinearProj(
      bias=f32[1024],
      weight=f32[2048,1024],
      input_dim=2048,
      output_dim=1024,
      use_bias=True
    )
  ),
  post_ln=LayerNorm(
    shape=(1024,),
    eps=1e-05,
    use_weight=True,
    use_bias=True,
    weight=f32[1024],
    bias=f32[1024]
  ),
  unemb_ln=LayerNorm(
    shape=(1024,),
    eps=1e-05,
    use_weight=True,
    use_bias=True,
    weight=f32[1024],
    bias=f32[1024]
  ),
  out_head=LinearProj(
    bias=f32[50304],
    weight=f32[1024,50304],
    input_dim=1024,
    output_dim=50304,
    use_bias=True
  )
)