[benchmarks] Set autocast kwargs only if AMP.

[ysiraichi]

This PR sets the dtype kwarg only if AMP is supposed to be used. Previously, leaving the dtype kwarg made it so nullcontext complained about it. This would crash every inference benchmark.

cc @miladm

[frgossen]

Can you give a bit more detail on why this is needed?

[ysiraichi]

Otherwise, we are calling nullcontext(dtype=...), here:

xla/benchmarks/benchmark_model.py

Lines 155 to 159 in 3e5ca3c

	def train(self, inputs, collect_full_output=False):
	self._optimizer_zero_grad()
	with self.autocast(**self.autocast_kwargs):
	pred = self.module(*inputs)
	loss = self.compute_loss(pred)

Which ends up raising an error.

[frgossen]

Thanks!

[vanbasten23]

Otherwise, we are calling nullcontext(dtype=...), here:

xla/benchmarks/benchmark_model.py

Lines 155 to 159 in 3e5ca3c

	def train(self, inputs, collect_full_output=False):
	self._optimizer_zero_grad()
	with self.autocast(**self.autocast_kwargs):
	pred = self.module(*inputs)
	loss = self.compute_loss(pred)

Which ends up raising an error.

Not sure if I follow. With the change, for amp case, we'll call nullcontext(dtype=...) with non-empy kwargs.., wouldn't nullcontext complain too?

[ysiraichi]

As far as I understand it, for AMP case, self.autocast will be either torch.amp.autocast or torch.cuda.autocast. Am I missing something?