Enable fp8_weight recomputation during backwards pass

[drisspg]

Summary

Currently for both delayed and dynamic linear we have made a memory versus computation trade off. We always save the fp8_weight casted version of the weight for backwards.

On single node this will cause the both the model.weight tensors in high precision and the casted Float8Tensor version of the weight to exist on gpu memory until the casted weight is freed during backward pass.

We would like to provide an option to disable this.

This becomes more of an issue for existing FSDP implementation. The high precision wight will be sharded among the Gpus. When running forward, we will all-gather and unshard the weights to each gpu. But since since the nn.parameter is not used for mamtul( rather the casted float8 tensor is) the existing fsdp mechanism for removing the un-sharded weight will not work and we will end up saving the un-sharded tensor for backwards on each device. This extra memory usage scales w/ num gpus.

[drisspg]

Update here with two attempts to solve this:

• Add option for recomputing the casted weight during backwards #186
• Checkpoint to reduce fp8_weight tensor saved for backwards #193

[vkuzo]

pytorch/ao#562