Issue with too much memory usage when running in parallel with mx.distributed

[sck-at-ucy]

I am trying to refactor my physics-informed transformer model to run parallel. I am trying to create the various dataset arrays on rank=0, along also empty (zero) counterpart arrays of the same shape on rank > 0 and then use mx.distrubtred.all_sum to get copies of the dataset on all ranks. Then I slice the dataset arrays to create local copies for each rank. Once I no longer need the original complete dataset and I thought the code below would free the memory, but apparently is not, as the code uses way too much memory compared to running on a single node and the entire dataset.

What am I doing wrong? Is there a better memory management approach when using mx.distributed() in the absence of a distributed broadcast?

   local_training_data = training_data_mlx[start_train_idx:end_train_idx,:,:,:]
   local_training_alphas = training_alphas_mlx[start_train_idx:end_train_idx]
   local_training_dts = training_dts_mlx[start_train_idx:end_train_idx]
   #
   local_validation_data = validation_data_mlx[start_valid_idx:end_valid_idx,:,:,:]
   local_validation_alphas = validation_alphas_mlx[start_valid_idx:end_valid_idx]
   local_validation_dts = validation_dts_mlx[start_valid_idx:end_valid_idx]
   #
   local_test_data = test_data_mlx[start_test_idx:end_test_idx,:,:,:]
   local_test_alphas = test_alphas_mlx[start_test_idx:end_test_idx]
   local_test_dts = test_dts_mlx[start_test_idx:end_test_idx
   #
   mx.eval(local_training_data)
   mx.eval(local_validation_data)
   mx.eval(local_test_data)
    #
   mx.eval(local_training_alphas)
   mx.eval(local_validation_alphas)
   mx.eval(local_test_alphas)
   #
   mx.eval(local_training_dts)
   mx.eval(local_validation_dts)
   mx.eval(local_test_dts)
   #
   training_data_mlx = None
   validation_data_mlx = None
   test_data_mlx = None
    #
   training_alphas_mlx = None
   validation_alphas_mlx = None
   test_alphas_mlx = None
    #
   training_dts_mlx = None
   validation_dts_mlx = None
   test_dts_mlx = None

   mx.metal.clear_cache()

Would be thankful for any advice from @awni @angeloskath :)

[awni]

• Setting the arrays to None and clearing the cache should free the memory they used.

• Probably we should have a broadcast to avoid doing what you are doing which seems tedious, but that's a slightly different matter

• How are you measuring memory? And how much more is the multi-process version using than the single process version?

• Could you provide the full code or a simplified way to reproduce what you are seeing?