Default to allow_tf32=True for GPU Devices

[ravi-mosaicml]

PyTorch 1.12 disabled TF32 format by default for matmuls. This could lead to a significant performance regression.

To fix, adding a flag in the DeviceGPU class to control whether to set PyTorch's allow_tf32 flag. By default, Composer will set this to True, for consistent behavior across pytorch versions.

Closes https://mosaicml.atlassian.net/browse/CO-328

[abhi-mosaic]

Looks good to me! For posterity, I want to note that we put this flag in the DeviceGPU rather than as a Trainer precision enum, because it reflects how certain devices (NVIDIA Ampere+ GPUs ) perform "single precision" GEMM ops, rather than an overall training loop strategy for training with a particular precision.

At no point will any tensor or checkpoint or activation ever be stored in TF32, the format only exists inside NVIDIA tensor cores right before the inner products of a GEMM are computed, and then the accumulation and output tensor are true FP32. See this article

An analogy would be if we implemented DeviceTPU, we might want an allow_bf16=True flag which is forced to True, which is not the same as BF16 mixed precision training. In their default behavior, TPUs cast FP32 tensors to BF16 right before the GEMM inner products, and output FP32. The overall training strategy of BF16 mixed precision (with BF16 activations and gradients) is a separate step. See this article

[linden-li]

LGTM!