[benchmarks] PyTorch HUD runs on different data-types.

[ysiraichi]

Problem

Currently, we only change the data-type of the models on CUDA, and only if they have DEFAULT_CUDA_<test>_PRECISION specified. Otherwise, the models run on float32 precision. Meanwhile, PyTorch HUD runs inference on bfloat16 and training on AMP.

The data-type of the models is relevant not only for performance, but also for coverage. That is because depending on the data-type of the model, the amount of memory it uses differ.

Possible Solutions

In order to better compare the results on PyTorch HUD and the results we get from running the scripts that live in PyTorch/XLA repository, I think there are a couple of options:

• Introduce --bfloat16 and --amp arguments for forcing the models to be of a specific precision
• Introduce --pytorch-hud argument for the behavior described above
• Make PyTorch HUD used data-types be the default, and add --default-data-type for the old behavior

cc @miladm @golechwierowicz @cota @frgossen @zpcore @vanbasten23

[lezcano]

I think defaulting to the PyTorch HUD behaviour and moving the current behaviour under a flag makes sense. cc @miladm

[lezcano]

Perhaps we may not even want to keep the previous behaviour. I'll defer to google engs to decide this.

[frgossen]

Agreed, resembling PyTorch HUD behaviour as closely as possible makes sense.

[ManfeiBai]

Hi, @lsy323, is it ok to assign this ticket to you?

[ysiraichi]

@ManfeiBai @lsy323 Ah, sorry. I should have mentioned it. I'm working on this one.