I am confused by the function _balance_batch in RayPPOTrainer. I roughly understand what this function is supposed to do through its comment 'Reorder the data on single controller such that each dp rank gets similar total tokens'

So why passing world_size to the inner function get_seqlen_balanced_partitions rather than dp_size if it is supposed to do load balance across dp ranks ?