Sample fine-tuning environment

[wesselb]

• Removes bf16_mode and brings back autocast (Closes Illegal Memory Access with Mixed Precision #121.).
• Make activation checkpointing configurable and improve default checkpointing strategy.
• Change default drop-out rates to zero everywhere
• Fix bug: Set attn_drop_rate in the backbone. This will cause stochasticity if the model wasn't .eval()d.
• Adds a Docker image which runs a very basic fine-tuning loop for Aurora. (Closes Multi-GPU training and mixed precision #125. Closes Illegal Memory Access with Mixed Precision #121.)

I've verified that the instructions work on a fresh Azure VM.

[rchan26]

Thanks @wesselb - this looks great! Just some small comments. The main one is about ensuring that we definitely use the right .venv when in the image. Once we've made sure that is the case, I think we can merge this in :)

[rchan26]

looks good to me! thanks again @wesselb !

[wesselb]

Thanks for the review, @rchan26 :)