dtype conversion based on batch_size?!!

[mashaan14]

Hi everyone,

I load CIFAR-10 images and put them in a grain.DataLoader. I used a jax.sharding.Mesh:

mesh = Mesh(mesh_utils.create_device_mesh((8, 1)), ('batch', 'model'))

My model is a single convolutional layer that cuts images into patches:

class PatchEmbedding(nnx.Module):
    def __init__(self, img_size, patch_size, embed_dim, dtype, param_dtype, rngs: nnx.Rngs(0)):
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        fh, fw = patch_size, patch_size # Filter height/width are your patch_size

        self.conv_proj = nnx.Conv(
            in_features=3,
            out_features=embed_dim,          
            kernel_size=(fh, fw),        
            strides=(fh, fw),            
            padding='VALID',             
            kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), NamedSharding(mesh, P(None, 'model'))),
            bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), NamedSharding(mesh, P('model'))),
            dtype=dtype,                 
            # param_dtype=param_dtype,           
            rngs=rngs,                   
        )
        self.dtype = dtype

    def __call__(self, x: jax.Array) -> jax.Array:
        x = x.astype(self.dtype)
        x = self.conv_proj(x)

        n, h, w, c = x.shape
        x = jnp.reshape(x, [n, h * w, c])
        return x

and run a single batch:

model = PatchEmbedding(data_config.img_size,
                             model_config.patch_size,
                             model_config.embed_dim,
                             model_config.dtype,
                             model_config.param_dtype,
                             model_config.rngs)

log_dir = "./jax_profile_logs"
jax.profiler.start_trace(log_dir)

for batch in train_loader:
  images = jax.device_put(batch[0], NamedSharding(mesh, P('batch', None)))
  model(images)
  break

jax.profiler.stop_trace()

Here's my question:
when I used batch_size=128 I got a copy HLO OP that changes the dtype from bf16 to f32:

%copy = bf16[16,32,32,3]{0,3,2,1:T(4,128)(2,1)} copy(f32[16,32,32,3]{2,1,3,0:T(8,128)} %param), sharding={devices=[8,1,1,1]<=[8]}

%fusion = f32[16,8,8,384]{3,0,2,1:T(8,128)} fusion(bf16[16,32,32,3]{0,3,2,1:T(4,128)(2,1)} %copy, f32[4,4,3,384]{3,1,2,0:T(4,128)} %param.1), kind=kOutput, calls=%fused_computation

but when I use batch_size=4096 I don't get that copy operation

%fusion = f32[512,8,8,384]{3,0,2,1:T(8,128)} fusion(f32[512,32,32,3]{0,2,3,1:T(8,128)} %param, f32[4,4,3,384]{3,1,2,0:T(4,128)} %param.1), kind=kOutput, calls=%fused_computation

Does this copy operation occur based on the batch size? or is it something else that causes the dtype to change?

appreciate your help :)