[Misc] ThroughputHook
- Dominant language
- Python
- Stars
- 5.2k
- Forks
- 448
- Avg merge
- 3d 15h
- Merged PRs (30d)
- 26
Description
[throughput_hook.py](https://github.com/InternLM/xtuner/blob/main/xtuner/engine/hooks/throughput_hook.py) 中仅仅计算一个`micro batch size`的`tgs`,能否添加一个`global batch size` 的 `tgs`输出
当`accumulative_counts > 1`时,在最后一个梯度累计`iter`,由于比其他`iter`多一个`optim.step()`操作,因此直接通过`micro batch size`输出的`tgs`求均值,会导致比实际`tgs`大,尤其在`optim offload`并且`accumulative_counts`较小时
例如考虑单机单卡情况,`accumulative_counts`为`2`时,假设`batch size`为`1`,`sequence_len`为`s`,第一个`iter tgs`为 $\frac{s}{t_1}$,第二个`iter tgs`为 $\frac{s}{t_2}$ ,如果直接计算两个`iter tgs`均值,那么`gbs tgs`为 $\frac{(\frac{s}{t_1} + \frac{s}{t_2})}{2} = \frac{s (t_1+ t_2)}{2t_1t_2}$
但是实际`gbs tgs`计算应为 $\frac{2s}{t_1 + t_2}$
两者相除为 $\frac{(t_1+t_2)^2}{4t_1t_2} \geqslant 1$
当`sequence_len`固定时,通过在[throughput_hook.py](https://github.com/InternLM/xtuner/blob/main/xtuner/engine/hooks/throughput_hook.py) 中添加如下代码计算`global batch size tgs`
```
if (batch_idx+1) % runner.strategy.config['gradient_accumulation_steps'] == 0:
message_hub.update_scalar('train/gbs_tokens_per_sec',
batch_size * sequence_len / (
message_hub.get_scalar('train/time').mean(runner.strategy.config['gradient_accumulation_steps']) + 1e-12))
```
Contributor guide
Assessment
This issue has not been assessed yet.