[FEA] Refine training benchmark model configuration.

[JacoCheung]

Currently, our benchmark configuration can not reflect real scenarios well.

1. The batchsize or sequence length is not typical
2. Input data is uniform, we need powerlaw. (both seqlen and key range)
3. The embedding dim is equal to hidden dim
4. The attention hyper-params is not well set
5. The model weight is not large enough. (because of 3. and 4.)
6. We have only single node test.
7. We do not benchmark the TP.
   Therefore we need to improve our benchmark suites.

---

By submitting this issue, you agree to follow our code of conduct and our contributing guidelines.

[JacoCheung]

Hi @shijieliu . I would like issue #117 to address hstu layer perf only. We can move to this issue regarding model perf benchmark.

usually should be 1k - 8k

I think batchsize might be coupled with sequence length. If we use larger hidden_size and batchsize X seqlen. It may lead to OOM easily. In that sense, I think we need to add large scaling (like 16x8 GPUs) benchmark (or with real datasets).

I think we can use larger hidden_size
You mean we can use more neads or larger head_dim?

Does this benchmark reflect embedding as well? (i.e. use small dense but intensive dyanmic embedding).

[JacoCheung]

Update with more specfic requirements.

1. Setup a baseline: Sufficient embeddings + dense (at least 2 nodes)
   1. Embedding all offloaded to CPU.
   2. (need to balance the workload ration between sparse vs dense)
   3. Random dataset on CPU.
2. Turn on op level opt.
   1. HSTU cutlass
   2. Embedding Cache
   3. recompute
   4. TP
3. Pipeline (Prefetch)
4. Workload balancer

[shijieliu]

junjie is working on figuring out the suitable workload(seqlen distribution) and after verifying the optimization effect, will dig into how to hide workload balancer with prefetch.