Scan over layers with sharded parameters increases memory usage.

[AllanYangZhou]

I'm having a problem where scanning over layers whose parameters are sharded causes OOM. My naive implementation is for the forward pass to loop over layers, but JIT-ing the forward takes a long time. Scanning over layers is recommended to improve compilation time, but seems to increase memory usage.

Below is a minimal repro. If LOOP==True, the parameters are stored as num_layers separate arrays and the forward pass is computed using a for loop (my naive implementation). If LOOP==False, the parameters are stored in a single array and the forward pass is computed via a scan operation. On a TPU v3-8, the former works fine but the latter. OOMs.

from functools import partial
import jax
import jax.numpy as jnp
from jax.experimental import mesh_utils

Mesh, NamedSharding = jax.sharding.Mesh, jax.sharding.NamedSharding
P, with_sharding_constraint = jax.sharding.PartitionSpec, jax.lax.with_sharding_constraint

LOOP = True  # loop vs scan

batch, t = 64, 1024
num_layers, num_heads, head_size, embed_size  = 48, 48, 128, 9000

mesh = Mesh(mesh_utils.create_device_mesh((jax.device_count(),)), axis_names=('data',))
print(mesh)

if LOOP:
    qkv_sharding = [NamedSharding(mesh, P(None, None, 'data')) for _ in range(num_layers)]
    o_sharding = [NamedSharding(mesh, P(None, None, 'data')) for _ in range(num_layers)]
else:
    qkv_sharding = NamedSharding(mesh, P(None, None, None, 'data'))
    o_sharding = NamedSharding(mesh, P(None, None, None, 'data'))
x_sharding = NamedSharding(mesh, P('data', None, None))


@partial(jax.jit, out_shardings=(qkv_sharding, o_sharding, x_sharding))
def init():
    if LOOP:
        qkv = [jnp.ones((num_heads, head_size, embed_size), dtype = jnp.bfloat16) for _ in range(num_layers)]
        o = [jnp.ones((num_heads, head_size, embed_size), dtype=jnp.bfloat16) for _ in range(num_layers)]
    else:
        qkv = jnp.ones((num_layers, num_heads, head_size, embed_size), dtype = jnp.bfloat16)
        o = jnp.ones((num_layers, num_heads, head_size, embed_size), dtype=jnp.bfloat16)
    x = jnp.ones((batch, t, embed_size), dtype=jnp.bfloat16)
    return qkv, o, x


if LOOP:
    def forward(params, x):
        qkvs, os = params
        for qkv, o in zip(qkvs, os):
            y = jnp.einsum('bte,hde->bthd', x, qkv) 
            x = jnp.einsum('bthd,hde->bte', y, o)
        return jnp.mean(x**2)
else:
    def forward(params, x):
        def layer(_x, params):
            _qkv, _o = params
            y = jnp.einsum('bte,hde->bthd', _x, _qkv) 
            z = jnp.einsum('bthd,hde->bte', y, _o)
            return z, None
        out, _ = jax.lax.scan(layer, x, params)
        return jnp.mean(out**2)

key = jax.random.PRNGKey(0)
qkv, o, x = init()

backward = jax.jit(jax.grad(forward), out_shardings=(qkv_sharding, o_sharding))
params = (qkv, o)
out = backward(params, x)

Here are the biggest allocation:

  Largest program allocations in hbm:                                                                                             
                                                                                                                                  
  1. Size: 6.59G                                                                                                                  
     Operator: op_name="jit(forward)/jit(main)/broadcast_in_dim[shape=(48, 64, 1024, 9000) broadcast_dimensions=()]" source_file="
/home/allanzhou/midgpt/debug.py" source_line=64                                                                                   
     Shape: bf16[48,8,1024,9000]{2,3,1,0:T(8,128)(2,1)}                                                                           
     Unpadded size: 6.59G                                                                                                         
     XLA label: broadcast.10 = broadcast(constant.129), dimensions={}                                                             
     Allocation type: HLO temp                                                                                                    
     ==========================                                                                                                   
                                                                                                                                  
  2. Size: 4.50G                                                                                                                  
     Operator: op_name="jit(forward)/jit(main)/broadcast_in_dim[shape=(48, 64, 1024, 48, 128) broadcast_dimensions=()]" source_fil
e="/home/allanzhou/midgpt/debug.py" source_line=64                                                                                
     Shape: bf16[48,8,1024,48,128]{4,2,3,1,0:T(8,128)(2,1)}                                                                       
     Unpadded size: 4.50G                                                                                                         
     XLA label: broadcast.12 = broadcast(constant.129), dimensions={}                                                             
     Allocation type: HLO temp                                                                                                    
     ==========================                                              

Is there a way to prevent this or debug what's going on?

[mattjj]

Thanks for the question!

This is just a guess, but as mentioned in the practical notes for jax.remat, using scan can defeat some XLA rematerialization optimizations and thus cause OOMs. Maybe try applying jax.remat to the layer function, and see how it affects performance? You can use the policies described in those docs to control things more precisely.

[AllanYangZhou]

Thanks, I think that might be it! Adding jax.checkpoint to the scanned function removes the OOMs. Though, in my actual training loop it seems to be ~25% slower than without scan, and the predefined policies don't seem to help. I might have to figure out a custom policy or try recursive checkpointing to remove that discrepancy.