[FEA] Balanced jagged compute&memory among DP ranks

[JacoCheung]

The input of HSTU are inherently varied-length sequences (tensor), and thus dp batches can have different memory and compute resources demand. We need a way to balance the workload that takes both compute and memory resource into account.

[JacoCheung]

Refering veRL .
The method they adopt is on the premise on the standard attention ops (pure computations) . See code which include proj q,k,v and qk,kv batched gemm. And partition the sequences with a karmarkar_karp algorithm.

This is a rather ideal assumption.(compute-bound). We can elaborate it with memory-bound region. Or we can do run a few warmup with sampling sequence length to get more reliable latency/throughput.

examples/commons/perf_models/
... task_estimator.py: gemm and attention variations latency/throughput estimator and with optionally warmup runs (like the autotuning process)
... mods: GPU specs
... partitoner.py: a karmarkar_karp wrapper that accepts `workloads` and partition it as even as possible.

examples/commons: 
... batch_shuffler.py: comm prims. (we need to A2A the sequence info and reorder the batch across dp ranks. To achieve this, we need to well define a batch. )

[JacoCheung]

Had almost all the functions ready. But now having trouble in incorporating into train pipeline: While enabling balanced shuffle, I can see warning like : (movielen_ranking)

/usr/local/lib/python3.12/dist-packages/torchmetrics/utilities/prints.py:42: UserWarning: No positive samples in targets, true positive value should be meaningless. Returning zero tensor in true positive score

Which indicated that the label of some class was missing in the whole eval predictions. But I checked the internal labels with torch.unique and found out no missing. The auc looked promising.

[JacoCheung]

Besides, there's one action remained:

1. Enable in sid_gr.