Support dp_size in replay buffer

[DNXie]

Added a dp_size dimension in replay buffer sampling to enable data parallel.

Updated GRPO/main.py accordingly

Test:

1. pytest tests/unit_tests/test_replay_buffer.py

 7 passed, 7 warnings in 23.35s

python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml

All services initialized successfully!
Starting GRPO training loops...
Generated 10 rollouts w/ average reward 0.275
Completed 10 training steps
Latest loss: 0.01575610041618347
Generated 20 rollouts w/ average reward 0.05
Completed 20 training steps
Latest loss: 0.0045821815729141235

[pbontrager]

Looks good, just need to fix one thing to not return sorted samples

[pbontrager]

We don't want to return a sorted sample here as that reduces variability in the sample. You need to get the index of the sorted array and then probably do this as a nested for loop to be easier to read.

batch = []
for rank in self.dp_size:
	local_batch = []
	for i in bsz:
		e = sampled_episodes[sort_order[rank*i]]
		local_batch.append(e)
	batch.append(local_batch)

[DNXie]

Thanks for pointing out this issue. I have updated this part. Please review.

[pbontrager]

This looks good now. There's a few more small things to fix but I'll pre-approve it

[pbontrager]

This is cleaner but is it moving the data twice? It's probably fine.

[DNXie]

I have updated the code to make it more efficient.