使用多卡测试，怎样让显存分配更均匀？

[rationalspark]

需测试Qwen2-7b的LongBench性能，上下文长度要达到32k。服务器有8块A100 80G。
在模型的配置中设置run_cfg=dict(num_gpus=8, num_procs=1)
但是，实际运行时，发现显存分配不均，前面上下文长度较短时，两块卡60多G，其他就不到10G。到一个较长的数据集，就直接爆显存了。
请各位大牛指导一下，怎样设置才能让显存分配更平均？实在无法在文档中找到，只好请大家帮忙。
不胜感谢。